docfast/src/__tests__/usage-flush.test.ts
Hoid d2f819de94
All checks were successful
Build & Deploy to Staging / Build & Deploy to Staging (push) Successful in 12m5s
fix: flush usage entries independently to prevent batch poisoning (BUG-100)
2026-03-04 14:04:53 +01:00

57 lines
1.9 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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);
});
});