280 lines
11 KiB
TypeScript
280 lines
11 KiB
TypeScript
/**
|
||
* tests/unit/quota-sqlite-store.test.ts
|
||
*
|
||
* Coverage for src/lib/quota/sqliteQuotaStore.ts:
|
||
* - Happy path: consume + peek returns correct value
|
||
* - Two consecutive consumes → sum
|
||
* - Bucket rotation: decayed sliding window
|
||
* - Concurrency: 50 parallel consumes → exact sum (mutex guards)
|
||
* - poolUsageWithDimensions: validates shape of PoolUsageSnapshot
|
||
* - clear() zeroes consumption
|
||
*/
|
||
|
||
import test from "node:test";
|
||
import assert from "node:assert/strict";
|
||
import fs from "node:fs";
|
||
import os from "node:os";
|
||
import path from "node:path";
|
||
|
||
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-sqlite-store-"));
|
||
process.env.DATA_DIR = TEST_DATA_DIR;
|
||
|
||
const core = await import("../../src/lib/db/core.ts");
|
||
const poolsDb = await import("../../src/lib/db/quotaPools.ts");
|
||
const sqliteQuotaStoreModule = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
|
||
async function resetStorage() {
|
||
core.resetDbInstance();
|
||
for (let attempt = 0; attempt < 10; attempt++) {
|
||
try {
|
||
if (fs.existsSync(TEST_DATA_DIR)) {
|
||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||
}
|
||
break;
|
||
} catch (err: unknown) {
|
||
const e = err as { code?: string };
|
||
if ((e?.code === "EBUSY" || e?.code === "EPERM") && attempt < 9) {
|
||
await new Promise((resolve) => setTimeout(resolve, 50 * (attempt + 1)));
|
||
} else {
|
||
throw err;
|
||
}
|
||
}
|
||
}
|
||
fs.mkdirSync(TEST_DATA_DIR, { recursive: true });
|
||
}
|
||
|
||
test.beforeEach(async () => {
|
||
await resetStorage();
|
||
});
|
||
|
||
test.after(async () => {
|
||
core.resetDbInstance();
|
||
if (fs.existsSync(TEST_DATA_DIR)) {
|
||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||
}
|
||
});
|
||
|
||
test("sqliteQuotaStore public surface excludes removed singleton reset helper", () => {
|
||
assert.equal(Object.hasOwn(sqliteQuotaStoreModule, "resetSqliteQuotaStore"), false);
|
||
assert.equal(typeof sqliteQuotaStoreModule.getSqliteQuotaStore, "function");
|
||
assert.equal(typeof sqliteQuotaStoreModule.SqliteQuotaStore, "function");
|
||
});
|
||
|
||
// Helper: make a dimension key
|
||
function makeDim(poolId = "pool-test", unit = "tokens" as const, window = "hourly" as const) {
|
||
return { poolId, unit, window };
|
||
}
|
||
|
||
// ─── Happy path ──────────────────────────────────────────────────────────────
|
||
|
||
test("sqliteQuotaStore: consume(100) then peek returns ~100", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim();
|
||
|
||
await store.consume("key-1", dim, 100);
|
||
const effective = await store.peek("key-1", dim);
|
||
|
||
// Since both consume and peek happen in the same bucket (milliseconds apart),
|
||
// prev=0, elapsed≈0 → effective ≈ 100. Allow small delta for timing.
|
||
assert.ok(effective > 99, `Expected >99, got ${effective}`);
|
||
assert.ok(effective <= 100, `Expected <=100, got ${effective}`);
|
||
});
|
||
|
||
test("sqliteQuotaStore: peek on fresh key returns 0", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim();
|
||
|
||
const effective = await store.peek("key-never-consumed", dim);
|
||
assert.equal(effective, 0);
|
||
});
|
||
|
||
// ─── Two consecutive consumes ────────────────────────────────────────────────
|
||
|
||
test("sqliteQuotaStore: two consumes sum correctly", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim();
|
||
|
||
await store.consume("key-2", dim, 100);
|
||
await store.consume("key-2", dim, 200);
|
||
const effective = await store.peek("key-2", dim);
|
||
|
||
// 300 in current bucket, prev=0, elapsed≈0 → effective≈300
|
||
assert.ok(effective > 299, `Expected >299, got ${effective}`);
|
||
assert.ok(effective <= 300, `Expected <=300, got ${effective}`);
|
||
});
|
||
|
||
// ─── Bucket rotation and decayed sliding window ──────────────────────────────
|
||
|
||
test("sqliteQuotaStore: bucket rotation applies decay from prev bucket", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const { WINDOW_MS } = await import("../../src/lib/quota/dimensions.ts");
|
||
const { incrementBucket } = await import("../../src/lib/db/quotaConsumption.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim("pool-rotate", "tokens", "hourly");
|
||
const windowMs = WINDOW_MS["hourly"]; // 3600000 ms
|
||
|
||
// Simulate: prev bucket has 1000 tokens
|
||
const nowMs = Date.now();
|
||
const currentBucket = Math.floor(nowMs / windowMs);
|
||
const prevBucket = currentBucket - 1;
|
||
const dimKey = `pool-rotate:tokens:hourly`;
|
||
|
||
// Write directly to prev bucket (bypassing store)
|
||
incrementBucket("key-rotate", dimKey, prevBucket, 1000, nowMs - windowMs);
|
||
|
||
// Peek at 50% elapsed through current bucket
|
||
// We can't easily fake time without mocking Date.now, so we verify the formula
|
||
// by reading the pair directly and computing manually.
|
||
const { getPair } = await import("../../src/lib/db/quotaConsumption.ts");
|
||
const { curr, prev } = getPair("key-rotate", dimKey, currentBucket);
|
||
|
||
assert.equal(curr, 0, "curr bucket should be empty");
|
||
assert.equal(prev, 1000, "prev bucket should have 1000");
|
||
|
||
// The sliding window formula: prev × (1 - elapsed/window) + curr
|
||
// When elapsed is small (just started current bucket), prev contributes a lot
|
||
const currentBucketStartMs = currentBucket * windowMs;
|
||
const elapsed = nowMs - currentBucketStartMs;
|
||
const expectedEffective = 1000 * (1 - elapsed / windowMs) + 0;
|
||
|
||
const effective = await store.peek("key-rotate", dim);
|
||
// Allow ±1% tolerance for timing
|
||
const tolerance = expectedEffective * 0.01 + 1;
|
||
assert.ok(
|
||
Math.abs(effective - expectedEffective) < tolerance,
|
||
`Expected ≈${expectedEffective.toFixed(2)}, got ${effective.toFixed(2)}`
|
||
);
|
||
});
|
||
|
||
// ─── Concurrency: 50 parallel consumes ──────────────────────────────────────
|
||
|
||
test("sqliteQuotaStore: 50 concurrent consumes → exact sum (mutex guards)", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim("pool-concurrent", "tokens", "hourly");
|
||
|
||
const N = 50;
|
||
const COST = 10;
|
||
|
||
// Fire 50 concurrent consumes
|
||
await Promise.all(Array.from({ length: N }, () => store.consume("key-concurrent", dim, COST)));
|
||
|
||
const effective = await store.peek("key-concurrent", dim);
|
||
|
||
// Total should be exactly N × COST = 500 (within the same bucket)
|
||
const expected = N * COST;
|
||
// Allow ±0.1% for floating point
|
||
assert.ok(
|
||
Math.abs(effective - expected) < expected * 0.001 + 0.1,
|
||
`Expected ≈${expected}, got ${effective}`
|
||
);
|
||
});
|
||
|
||
// ─── poolUsageWithDimensions ─────────────────────────────────────────────────
|
||
|
||
test("sqliteQuotaStore: poolUsageWithDimensions returns correct shape", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
|
||
// Create a real pool with allocations
|
||
const pool = poolsDb.createPool({
|
||
connectionId: "conn-pool-usage",
|
||
name: "Test Pool",
|
||
allocations: [
|
||
{ apiKeyId: "key-a", weight: 60, policy: "hard" },
|
||
{ apiKeyId: "key-b", weight: 40, policy: "soft" },
|
||
],
|
||
});
|
||
|
||
const dim = makeDim(pool.id, "tokens", "hourly");
|
||
await store.consume("key-a", dim, 300);
|
||
await store.consume("key-b", dim, 200);
|
||
|
||
const snapshot = await store.poolUsageWithDimensions(pool.id, [
|
||
{ unit: "tokens", window: "hourly", limit: 1000 },
|
||
]);
|
||
|
||
assert.equal(snapshot.poolId, pool.id);
|
||
assert.ok(snapshot.generatedAt, "generatedAt should be set");
|
||
assert.ok(Array.isArray(snapshot.dimensions), "dimensions should be array");
|
||
assert.equal(snapshot.dimensions.length, 1);
|
||
|
||
const dimSnap = snapshot.dimensions[0];
|
||
assert.equal(dimSnap.unit, "tokens");
|
||
assert.equal(dimSnap.window, "hourly");
|
||
assert.equal(dimSnap.limit, 1000);
|
||
// consumedTotal should be close to 300 + 200 = 500
|
||
assert.ok(
|
||
dimSnap.consumedTotal > 490,
|
||
`consumedTotal should be close to 500, got ${dimSnap.consumedTotal}`
|
||
);
|
||
assert.equal(dimSnap.perKey.length, 2);
|
||
|
||
// Validate perKey shapes
|
||
for (const pk of dimSnap.perKey) {
|
||
assert.ok(typeof pk.apiKeyId === "string");
|
||
assert.ok(typeof pk.consumed === "number");
|
||
assert.ok(typeof pk.fairShare === "number");
|
||
assert.ok(typeof pk.deficit === "number");
|
||
assert.ok(typeof pk.borrowing === "boolean");
|
||
}
|
||
});
|
||
|
||
test("sqliteQuotaStore: poolUsage for non-existent pool returns empty snapshot", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
|
||
const snapshot = await store.poolUsage("nonexistent-pool-id");
|
||
assert.equal(snapshot.poolId, "nonexistent-pool-id");
|
||
assert.equal(snapshot.dimensions.length, 0);
|
||
});
|
||
|
||
// ─── clear() ────────────────────────────────────────────────────────────────
|
||
|
||
test("sqliteQuotaStore: clear() zeroes consumption", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim("pool-clear", "tokens", "hourly");
|
||
|
||
await store.consume("key-clear", dim, 500);
|
||
const before = await store.peek("key-clear", dim);
|
||
assert.ok(before > 0, "Should have consumed some");
|
||
|
||
await store.clear("key-clear", dim);
|
||
const after = await store.peek("key-clear", dim);
|
||
// After clear, curr=0, prev=0 (both zeroed), so effective=0
|
||
assert.equal(after, 0);
|
||
});
|
||
|
||
test("sqliteQuotaStore: clear() on fresh key is a no-op (no error)", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim("pool-clear-noop", "tokens", "hourly");
|
||
|
||
// Should not throw
|
||
await store.clear("key-fresh", dim);
|
||
const val = await store.peek("key-fresh", dim);
|
||
assert.equal(val, 0);
|
||
});
|
||
|
||
// ─── Multiple keys, same dimension (isolation) ───────────────────────────────
|
||
|
||
test("sqliteQuotaStore: different keys are isolated", async () => {
|
||
const { SqliteQuotaStore } = await import("../../src/lib/quota/sqliteQuotaStore.ts");
|
||
const store = new SqliteQuotaStore();
|
||
const dim = makeDim("pool-iso", "tokens", "hourly");
|
||
|
||
await store.consume("key-iso-a", dim, 100);
|
||
await store.consume("key-iso-b", dim, 200);
|
||
|
||
const a = await store.peek("key-iso-a", dim);
|
||
const b = await store.peek("key-iso-b", dim);
|
||
|
||
// Each key should only see its own consumption
|
||
assert.ok(a < 110, `key-a should not see key-b's consumption, got ${a}`);
|
||
assert.ok(b > 190, `key-b should have its own consumption, got ${b}`);
|
||
});
|