fix: flush usage entries independently to prevent batch poisoning (BUG-100)
All checks were successful
Build & Deploy to Staging / Build & Deploy to Staging (push) Successful in 12m5s
All checks were successful
Build & Deploy to Staging / Build & Deploy to Staging (push) Successful in 12m5s
This commit is contained in:
parent
314edc182a
commit
d2f819de94
2 changed files with 81 additions and 32 deletions
57
src/__tests__/usage-flush.test.ts
Normal file
57
src/__tests__/usage-flush.test.ts
Normal file
|
|
@ -0,0 +1,57 @@
|
|||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
|
||||
// Unmock usage middleware — we want to test the real implementation
|
||||
vi.unmock("../middleware/usage.js");
|
||||
|
||||
import { connectWithRetry } from "../services/db.js";
|
||||
|
||||
describe("flushDirtyEntries – independent key flushing", () => {
|
||||
let usageMod: typeof import("../middleware/usage.js");
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.clearAllMocks();
|
||||
vi.resetModules();
|
||||
// Re-import to get fresh state
|
||||
usageMod = await import("../middleware/usage.js");
|
||||
});
|
||||
|
||||
it("should flush remaining keys even if one key fails", async () => {
|
||||
// Set up two keys in the in-memory cache via the middleware
|
||||
const next = vi.fn();
|
||||
const res = { status: vi.fn(() => ({ json: vi.fn() })) };
|
||||
|
||||
usageMod.usageMiddleware({ apiKeyInfo: { key: "key-good" } }, res, next);
|
||||
usageMod.usageMiddleware({ apiKeyInfo: { key: "key-bad" } }, res, next);
|
||||
|
||||
// Track which keys were successfully upserted
|
||||
const flushedKeys: string[] = [];
|
||||
let callCount = 0;
|
||||
|
||||
const mockQuery = vi.fn().mockImplementation((sql: string, params?: any[]) => {
|
||||
if (sql.includes("INSERT INTO usage")) {
|
||||
callCount++;
|
||||
if (params && params[0] === "key-bad") {
|
||||
throw new Error("simulated constraint violation");
|
||||
}
|
||||
flushedKeys.push(params![0]);
|
||||
}
|
||||
return { rows: [], rowCount: 0 };
|
||||
});
|
||||
|
||||
const mockRelease = vi.fn();
|
||||
|
||||
// Each call to connectWithRetry returns a fresh client
|
||||
vi.mocked(connectWithRetry).mockImplementation(async () => ({
|
||||
query: mockQuery,
|
||||
release: mockRelease,
|
||||
}) as any);
|
||||
|
||||
// Access the flush function (exported for testing)
|
||||
await usageMod.flushDirtyEntries();
|
||||
|
||||
// The good key should have been flushed despite the bad key failing
|
||||
expect(flushedKeys).toContain("key-good");
|
||||
// Release should be called for each key (independent clients)
|
||||
expect(mockRelease.mock.calls.length).toBeGreaterThanOrEqual(2);
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue