fix: tests

ratelimit/store/memory.nim

@@ -1,4 +1,4 @@
-import std/times
+import std/[times, options]
 import ./store
 import chronos
 
@@ -8,8 +8,6 @@ type MemoryRateLimitStore* = ref object
 
 proc newMemoryRateLimitStore*(): MemoryRateLimitStore =
   result = MemoryRateLimitStore()
-  result.bucketState =
-    BucketState(budget: 10, budgetCap: 10, lastTimeFull: Moment.now())
 
 proc saveBucketState*(
     store: MemoryRateLimitStore, bucketState: BucketState
@@ -17,5 +15,7 @@ proc saveBucketState*(
   store.bucketState = bucketState
   return true
 
-proc loadBucketState*(store: MemoryRateLimitStore): Future[BucketState] {.async.} =
-  return store.bucketState
+proc loadBucketState*(
+    store: MemoryRateLimitStore
+): Future[Option[BucketState]] {.async.} =
+  return some(store.bucketState)

ratelimit/store/sqlite.nim

@@ -1,6 +1,4 @@
-import std/times
-import std/strutils
-import std/json
+import std/[times, strutils, json, options]
 import ./store
 import chronos
 import db_connector/db_sqlite
@@ -12,47 +10,8 @@ type SqliteRateLimitStore* = ref object
 
 const BUCKET_STATE_KEY = "rate_limit_bucket_state"
 
-proc newSqliteRateLimitStore*(dbPath: string = ":memory:"): SqliteRateLimitStore =
-  result = SqliteRateLimitStore(dbPath: dbPath)
-  result.db = open(dbPath, "", "", "")
-
-  # Create table if it doesn't exist
-  result.db.exec(
-    sql"""
-    CREATE TABLE IF NOT EXISTS kv_store (
-        key TEXT PRIMARY KEY,
-        value BLOB
-    )
-  """
-  )
-
-  # Insert default state if key doesn't exist
-  let count = result.db
-    .getValue(sql"SELECT COUNT(*) FROM kv_store WHERE key = ?", BUCKET_STATE_KEY)
-    .parseInt()
-  if count == 0:
-    let defaultTimeSeconds = Moment.now().epochSeconds()
-    let defaultState = BucketState(
-      budget: 10,
-      budgetCap: 10,
-      lastTimeFull: Moment.init(defaultTimeSeconds, chronos.seconds(1)),
-    )
-
-    # Serialize to JSON
-    let jsonState =
-      %*{
-        "budget": defaultState.budget,
-        "budgetCap": defaultState.budgetCap,
-        "lastTimeFullSeconds": defaultTimeSeconds,
-      }
-
-    result.db.exec(
-      sql"INSERT INTO kv_store (key, value) VALUES (?, ?)", BUCKET_STATE_KEY, $jsonState
-    )
-
-proc close*(store: SqliteRateLimitStore) =
-  if store.db != nil:
-    store.db.close()
+proc newSqliteRateLimitStore*(db: DbConn): SqliteRateLimitStore =
+  result = SqliteRateLimitStore(db: db)
 
 proc saveBucketState*(
     store: SqliteRateLimitStore, bucketState: BucketState
@@ -61,32 +20,37 @@ proc saveBucketState*(
     # Convert Moment to Unix seconds for storage
     let lastTimeSeconds = bucketState.lastTimeFull.epochSeconds()
 
-    # Serialize to JSON
     let jsonState =
       %*{
         "budget": bucketState.budget,
         "budgetCap": bucketState.budgetCap,
         "lastTimeFullSeconds": lastTimeSeconds,
       }
-
     store.db.exec(
-      sql"UPDATE kv_store SET value = ? WHERE key = ?", $jsonState, BUCKET_STATE_KEY
+      sql"INSERT INTO kv_store (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value",
+      BUCKET_STATE_KEY,
+      $jsonState,
     )
     return true
   except:
     return false
 
-proc loadBucketState*(store: SqliteRateLimitStore): Future[BucketState] {.async.} =
+proc loadBucketState*(
+    store: SqliteRateLimitStore
+): Future[Option[BucketState]] {.async.} =
   let jsonStr =
     store.db.getValue(sql"SELECT value FROM kv_store WHERE key = ?", BUCKET_STATE_KEY)
+  if jsonStr == "":
+    return none(BucketState)
 
-  # Parse JSON and reconstruct BucketState
   let jsonData = parseJson(jsonStr)
   let unixSeconds = jsonData["lastTimeFullSeconds"].getInt().int64
   let lastTimeFull = Moment.init(unixSeconds, chronos.seconds(1))
 
-  return BucketState(
-    budget: jsonData["budget"].getInt(),
-    budgetCap: jsonData["budgetCap"].getInt(),
-    lastTimeFull: lastTimeFull,
+  return some(
+    BucketState(
+      budget: jsonData["budget"].getInt(),
+      budgetCap: jsonData["budgetCap"].getInt(),
+      lastTimeFull: lastTimeFull,
+    )
   )

ratelimit/store/store.nim

@@ -1,4 +1,4 @@
-import std/[times, deques]
+import std/[times, deques, options]
 import chronos
 
 type
@@ -10,4 +10,4 @@ type
   RateLimitStoreConcept* =
     concept s
         s.saveBucketState(BucketState) is Future[bool]
-        s.loadBucketState() is Future[BucketState]
+        s.loadBucketState() is Future[Option[BucketState]]

tests/test_sqlite_store.nim

@@ -2,28 +2,37 @@ import testutils/unittests
 import ../ratelimit/store/sqlite
 import ../ratelimit/store/store
 import chronos
+import db_connector/db_sqlite
+import ../chat_sdk/migration
+import std/[options, os]
 
 suite "SqliteRateLimitStore Tests":
-  asyncTest "newSqliteRateLimitStore - creates store with default values":
-    ## Given & When
-    let now = Moment.now()
-    let store = newSqliteRateLimitStore()
-    defer:
-      store.close()
+  setup:
+    let db = open("test-ratelimit.db", "", "", "")
+    runMigrations(db)
+
+  teardown:
+    if db != nil:
+      db.close()
+    if fileExists("test-ratelimit.db"):
+      removeFile("test-ratelimit.db")
+
+  asyncTest "newSqliteRateLimitStore - empty state":
+    ## Given
+    let store = newSqliteRateLimitStore(db)
+
+    ## When
+    let loadedState = await store.loadBucketState()
 
     ## Then
-    let bucketState = await store.loadBucketState()
-    check bucketState.budget == 10
-    check bucketState.budgetCap == 10
-    check bucketState.lastTimeFull.epochSeconds() == now.epochSeconds()
+    check loadedState.isNone()
 
   asyncTest "saveBucketState and loadBucketState - state persistence":
     ## Given
-    let now = Moment.now()
-    let store = newSqliteRateLimitStore()
-    defer:
-      store.close()
+    let store = newSqliteRateLimitStore(db)
 
+    let now = Moment.now()
+    echo "now: ", now.epochSeconds()
     let newBucketState = BucketState(budget: 5, budgetCap: 20, lastTimeFull: now)
 
     ## When
@@ -32,7 +41,8 @@ suite "SqliteRateLimitStore Tests":
 
     ## Then
     check saveResult == true
-    check loadedState.budget == newBucketState.budget
-    check loadedState.budgetCap == newBucketState.budgetCap
-    check loadedState.lastTimeFull.epochSeconds() ==
+    check loadedState.isSome()
+    check loadedState.get().budget == newBucketState.budget
+    check loadedState.get().budgetCap == newBucketState.budgetCap
+    check loadedState.get().lastTimeFull.epochSeconds() ==
       newBucketState.lastTimeFull.epochSeconds()