add unit tests for usage middleware (14 tests)

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