From 44707d9247e1ce5d05e8e9af53247b4c912a74c8 Mon Sep 17 00:00:00 2001
From: OpenClaw Subagent <subagent@openclaw.com>
Date: Fri, 13 Mar 2026 08:06:38 +0100
Subject: [PATCH] test: improve usage.ts coverage (getUsageForKey, retry
 exhaustion)

---
 src/__tests__/usage-coverage.test.ts | 113 +++++++++++++++++++++++++++
 1 file changed, 113 insertions(+)
 create mode 100644 src/__tests__/usage-coverage.test.ts

diff --git a/src/__tests__/usage-coverage.test.ts b/src/__tests__/usage-coverage.test.ts
new file mode 100644
index 0000000..19c56fe
--- /dev/null
+++ b/src/__tests__/usage-coverage.test.ts
@@ -0,0 +1,113 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+
+vi.unmock("../middleware/usage.js");
+
+import { connectWithRetry } from "../services/db.js";
+
+describe("usage.ts – coverage improvements", () => {
+  let usageMod: typeof import("../middleware/usage.js");
+  const next = vi.fn();
+  const res = { status: vi.fn(() => ({ json: vi.fn() })) } as any;
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    usageMod = await import("../middleware/usage.js");
+  });
+
+  describe("getUsageForKey", () => {
+    it("returns correct data for existing key in current month", () => {
+      // Track a key via middleware to populate in-memory cache
+      usageMod.usageMiddleware({ apiKeyInfo: { key: "test-key-123" } } as any, res, next);
+      usageMod.usageMiddleware({ apiKeyInfo: { key: "test-key-123" } } as any, res, next);
+
+      const result = usageMod.getUsageForKey("test-key-123");
+      expect(result.count).toBe(2);
+      expect(result.monthKey).toMatch(/^\d{4}-\d{2}$/);
+    });
+
+    it("returns {count: 0} for unknown key", () => {
+      const result = usageMod.getUsageForKey("nonexistent-key");
+      expect(result.count).toBe(0);
+      expect(result.monthKey).toMatch(/^\d{4}-\d{2}$/);
+    });
+
+    it("returns {count: 0} when month differs (stale data)", async () => {
+      // Track a key first
+      usageMod.usageMiddleware({ apiKeyInfo: { key: "stale-key" } } as any, res, next);
+
+      // Verify it exists
+      expect(usageMod.getUsageForKey("stale-key").count).toBe(1);
+
+      // Now mock Date to return a different month
+      const realDate = globalThis.Date;
+      const mockDate = class extends realDate {
+        constructor(...args: any[]) {
+          if (args.length === 0) {
+            super(2099, 11, 1); // Dec 2099 — different month
+          } else {
+            super(...(args as [any]));
+          }
+        }
+        static now() { return new realDate(2099, 11, 1).getTime(); }
+      };
+      vi.stubGlobal("Date", mockDate);
+
+      const result = usageMod.getUsageForKey("stale-key");
+      expect(result.count).toBe(0);
+      expect(result.monthKey).toBe("2099-12");
+
+      vi.stubGlobal("Date", realDate);
+    });
+  });
+
+  describe("flushDirtyEntries – retry exhaustion", () => {
+    it("removes key from dirty set after MAX_RETRIES failures", async () => {
+      // Track a key
+      usageMod.usageMiddleware({ apiKeyInfo: { key: "fail-key" } } as any, res, next);
+
+      const mockRelease = vi.fn();
+      vi.mocked(connectWithRetry).mockResolvedValue({
+        query: vi.fn().mockRejectedValue(new Error("db error")),
+        release: mockRelease,
+      } as any);
+
+      // Flush 3 times — each time it should retry, on 3rd it should give up
+      await usageMod.flushDirtyEntries(); // retry 1
+      await usageMod.flushDirtyEntries(); // retry 2
+      await usageMod.flushDirtyEntries(); // retry 3 → removed
+
+      // 4th flush should have nothing to do (key removed from dirty set)
+      mockRelease.mockClear();
+      await usageMod.flushDirtyEntries();
+      // connectWithRetry should not be called since dirtyKeys is empty
+      expect(mockRelease).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("threshold flush trigger", () => {
+    it("triggers flush when dirtyKeys.size >= FLUSH_THRESHOLD (50)", async () => {
+      const mockRelease = vi.fn();
+      const mockQuery = vi.fn().mockResolvedValue({ rows: [], rowCount: 0 });
+      vi.mocked(connectWithRetry).mockResolvedValue({
+        query: mockQuery,
+        release: mockRelease,
+      } as any);
+
+      // Track 50 unique keys to trigger threshold
+      for (let i = 0; i < 50; i++) {
+        usageMod.usageMiddleware(
+          { apiKeyInfo: { key: `threshold-key-${i}` } } as any,
+          res,
+          next,
+        );
+      }
+
+      // Give the async flush a tick to complete
+      await new Promise((r) => setTimeout(r, 50));
+
+      // Flush should have been triggered — connectWithRetry called for each key
+      expect(mockRelease.mock.calls.length).toBeGreaterThanOrEqual(50);
+    });
+  });
+});