From 9f52377d44c714e1b764a0ef624fb723bec40c3e Mon Sep 17 00:00:00 2001 From: pablo Date: Wed, 16 Jul 2025 10:05:47 +0300 Subject: [PATCH] fix: tests --- ratelimit/store/memory.nim | 10 +++--- ratelimit/store/sqlite.nim | 70 +++++++++---------------------------- ratelimit/store/store.nim | 4 +-- tests/test_sqlite_store.nim | 44 ++++++++++++++--------- 4 files changed, 51 insertions(+), 77 deletions(-) diff --git a/ratelimit/store/memory.nim b/ratelimit/store/memory.nim index ef19de1..557a17d 100644 --- a/ratelimit/store/memory.nim +++ b/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) diff --git a/ratelimit/store/sqlite.nim b/ratelimit/store/sqlite.nim index ba3e72e..e364e5d 100644 --- a/ratelimit/store/sqlite.nim +++ b/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, + ) ) diff --git a/ratelimit/store/store.nim b/ratelimit/store/store.nim index fab578d..c4f6da3 100644 --- a/ratelimit/store/store.nim +++ b/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]] diff --git a/tests/test_sqlite_store.nim b/tests/test_sqlite_store.nim index 8728992..6315ec9 100644 --- a/tests/test_sqlite_store.nim +++ b/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()