From c01e88686a1800a4ea8585e80c0ea5d2f91cad45 Mon Sep 17 00:00:00 2001 From: OpenClaw Date: Thu, 26 Feb 2026 13:04:15 +0000 Subject: [PATCH] add unit tests for usage middleware (14 tests) --- src/__tests__/usage.test.ts | 231 ++++++++++++++++++++++++++++++++++++ 1 file changed, 231 insertions(+) create mode 100644 src/__tests__/usage.test.ts diff --git a/src/__tests__/usage.test.ts b/src/__tests__/usage.test.ts new file mode 100644 index 0000000..2a1efbc --- /dev/null +++ b/src/__tests__/usage.test.ts @@ -0,0 +1,231 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; + +// Unmock usage middleware — we want to test the real implementation +vi.unmock("../middleware/usage.js"); + +// DB and keys are still mocked by setup.ts +import { queryWithRetry } from "../services/db.js"; +import { isProKey } from "../services/keys.js"; + +describe("usage middleware", () => { + let usage: typeof import("../middleware/usage.js"); + + const mockJson = vi.fn(); + const mockStatus = vi.fn(() => ({ json: mockJson })); + const mockNext = vi.fn(); + + function makeReq(key = "df_free_testkey123"): any { + return { apiKeyInfo: { key } }; + } + + function makeRes(): any { + mockJson.mockClear(); + mockStatus.mockClear(); + return { status: mockStatus, json: mockJson }; + } + + beforeEach(async () => { + vi.clearAllMocks(); + vi.resetModules(); + // Re-import to get fresh internal state (new Map) + usage = await import("../middleware/usage.js"); + }); + + // --- loadUsageData --- + + describe("loadUsageData", () => { + it("populates in-memory map from DB rows", async () => { + vi.mocked(queryWithRetry).mockResolvedValueOnce({ + rows: [ + { key: "df_free_abc12345", count: 42, month_key: "2026-02" }, + ], + rowCount: 1, + } as any); + + await usage.loadUsageData(); + + const stats = usage.getUsageStats("df_free_abc12345"); + expect(stats["df_free_..."]).toEqual({ count: 42, month: "2026-02" }); + }); + + it("handles empty DB result gracefully", async () => { + vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 0 } as any); + await usage.loadUsageData(); + const stats = usage.getUsageStats("df_free_nonexist"); + expect(stats).toEqual({}); + }); + + it("handles DB error gracefully (starts fresh)", async () => { + vi.mocked(queryWithRetry).mockRejectedValueOnce(new Error("DB down")); + await usage.loadUsageData(); + const stats = usage.getUsageStats("anything"); + expect(stats).toEqual({}); + }); + }); + + // --- getUsageStats --- + + describe("getUsageStats", () => { + it("returns empty object when key not found", () => { + expect(usage.getUsageStats("nonexistent")).toEqual({}); + }); + + it("returns empty object when no key provided", () => { + expect(usage.getUsageStats()).toEqual({}); + }); + + it("masks key to first 8 chars + '...'", async () => { + vi.mocked(queryWithRetry).mockResolvedValueOnce({ + rows: [{ key: "df_free_longkeyvalue", count: 5, month_key: "2026-02" }], + rowCount: 1, + } as any); + await usage.loadUsageData(); + + const stats = usage.getUsageStats("df_free_longkeyvalue"); + const keys = Object.keys(stats); + expect(keys).toHaveLength(1); + expect(keys[0]).toBe("df_free_..."); + }); + }); + + // --- usageMiddleware --- + + describe("usageMiddleware", () => { + it("calls next() for free key under limit", () => { + vi.mocked(isProKey).mockReturnValue(false); + const req = makeReq(); + const res = makeRes(); + + usage.usageMiddleware(req, res, mockNext); + + expect(mockNext).toHaveBeenCalled(); + expect(mockStatus).not.toHaveBeenCalled(); + }); + + it("calls next() for pro key under limit", () => { + vi.mocked(isProKey).mockReturnValue(true); + const req = makeReq("df_pro_testkey123"); + const res = makeRes(); + + usage.usageMiddleware(req, res, mockNext); + + expect(mockNext).toHaveBeenCalled(); + expect(mockStatus).not.toHaveBeenCalled(); + }); + + it("returns 429 when free key exceeds 100/month", async () => { + vi.mocked(isProKey).mockReturnValue(false); + const key = "df_free_limited1234"; + + // Seed usage at the limit + const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`; + vi.mocked(queryWithRetry).mockResolvedValueOnce({ + rows: [{ key, count: 100, month_key: monthKey }], + rowCount: 1, + } as any); + await usage.loadUsageData(); + + const req = makeReq(key); + const res = makeRes(); + + usage.usageMiddleware(req, res, mockNext); + + expect(mockStatus).toHaveBeenCalledWith(429); + expect(mockJson).toHaveBeenCalledWith( + expect.objectContaining({ error: expect.stringContaining("Free tier limit") }) + ); + expect(mockNext).not.toHaveBeenCalled(); + }); + + it("returns 429 when pro key exceeds 5000/month", async () => { + vi.mocked(isProKey).mockReturnValue(true); + const key = "df_pro_limited12345"; + + const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`; + vi.mocked(queryWithRetry).mockResolvedValueOnce({ + rows: [{ key, count: 5000, month_key: monthKey }], + rowCount: 1, + } as any); + await usage.loadUsageData(); + + const req = makeReq(key); + const res = makeRes(); + + usage.usageMiddleware(req, res, mockNext); + + expect(mockStatus).toHaveBeenCalledWith(429); + expect(mockJson).toHaveBeenCalledWith( + expect.objectContaining({ error: expect.stringContaining("Pro tier limit") }) + ); + expect(mockNext).not.toHaveBeenCalled(); + }); + + it("increments usage count on each call", () => { + vi.mocked(isProKey).mockReturnValue(false); + const key = "df_free_counting123"; + + // Call middleware 3 times + for (let i = 0; i < 3; i++) { + usage.usageMiddleware(makeReq(key), makeRes(), mockNext); + } + + const stats = usage.getUsageStats(key); + const masked = Object.keys(stats)[0]; + expect(stats[masked].count).toBe(3); + }); + + it("resets count when month changes", async () => { + vi.mocked(isProKey).mockReturnValue(false); + const key = "df_free_monthreset"; + + // Seed with old month data at limit + vi.mocked(queryWithRetry).mockResolvedValueOnce({ + rows: [{ key, count: 100, month_key: "2025-01" }], + rowCount: 1, + } as any); + await usage.loadUsageData(); + + // Current month is different, so it should reset and allow + const req = makeReq(key); + const res = makeRes(); + usage.usageMiddleware(req, res, mockNext); + + expect(mockNext).toHaveBeenCalled(); + expect(mockStatus).not.toHaveBeenCalled(); + + // Count should be 1 (reset + this request) + const stats = usage.getUsageStats(key); + const masked = Object.keys(stats)[0]; + expect(stats[masked].count).toBe(1); + }); + + it("allows free key at count 99 (just under limit)", async () => { + vi.mocked(isProKey).mockReturnValue(false); + const key = "df_free_almost1234"; + const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`; + + vi.mocked(queryWithRetry).mockResolvedValueOnce({ + rows: [{ key, count: 99, month_key: monthKey }], + rowCount: 1, + } as any); + await usage.loadUsageData(); + + const req = makeReq(key); + const res = makeRes(); + usage.usageMiddleware(req, res, mockNext); + + expect(mockNext).toHaveBeenCalled(); + expect(mockStatus).not.toHaveBeenCalled(); + }); + + it("handles missing apiKeyInfo gracefully (uses 'unknown')", () => { + vi.mocked(isProKey).mockReturnValue(false); + const req = { apiKeyInfo: undefined } as any; + const res = makeRes(); + + usage.usageMiddleware(req, res, mockNext); + + expect(mockNext).toHaveBeenCalled(); + }); + }); +});