From ae8b32e1c45125f294a4aaf1281c814929940f9e Mon Sep 17 00:00:00 2001
From: OpenClaw Subagent <subagent@openclaw.com>
Date: Thu, 12 Mar 2026 20:08:23 +0100
Subject: [PATCH] test: improve db.ts and keys.ts coverage

---
 src/__tests__/db-init-cleanup.test.ts |  99 ++++++++++++++++++++
 src/__tests__/keys-coverage.test.ts   | 129 ++++++++++++++++++++++++++
 2 files changed, 228 insertions(+)
 create mode 100644 src/__tests__/db-init-cleanup.test.ts
 create mode 100644 src/__tests__/keys-coverage.test.ts

diff --git a/src/__tests__/db-init-cleanup.test.ts b/src/__tests__/db-init-cleanup.test.ts
new file mode 100644
index 0000000..2421697
--- /dev/null
+++ b/src/__tests__/db-init-cleanup.test.ts
@@ -0,0 +1,99 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+
+// Mock pg and logger like db.test.ts does
+const mockRelease = vi.fn();
+const mockQuery = vi.fn();
+const mockConnect = vi.fn();
+
+vi.mock("pg", () => {
+  const Pool = vi.fn(function () {
+    return {
+      connect: mockConnect,
+      on: vi.fn(),
+    };
+  });
+  return { default: { Pool }, Pool };
+});
+
+vi.mock("../services/logger.js", () => ({
+  default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
+}));
+
+// Use real db.js implementation
+vi.mock("../services/db.js", async () => {
+  return await vi.importActual("../services/db.js");
+});
+
+describe("initDatabase", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockConnect.mockReset();
+    mockQuery.mockReset();
+    mockRelease.mockReset();
+  });
+
+  it("calls connectWithRetry, runs DDL, and releases client", async () => {
+    // connectWithRetry does pool.connect() then SELECT 1 validation
+    mockQuery.mockResolvedValue({ rows: [], rowCount: 0 });
+    mockConnect.mockResolvedValue({ query: mockQuery, release: mockRelease });
+
+    const { initDatabase } = await import("../services/db.js");
+    await initDatabase();
+
+    // SELECT 1 validation + DDL = at least 2 calls
+    expect(mockQuery).toHaveBeenCalledTimes(2);
+    // DDL should contain CREATE TABLE
+    const ddlCall = mockQuery.mock.calls[1][0] as string;
+    expect(ddlCall).toContain("CREATE TABLE IF NOT EXISTS api_keys");
+    expect(ddlCall).toContain("CREATE TABLE IF NOT EXISTS usage");
+    expect(mockRelease).toHaveBeenCalledWith();
+  });
+
+  it("releases client even if DDL fails", async () => {
+    const selectQuery = vi.fn()
+      .mockResolvedValueOnce({ rows: [{ "?column?": 1 }] }) // SELECT 1
+      .mockRejectedValueOnce(new Error("DDL failed")); // DDL
+    mockConnect.mockResolvedValue({ query: selectQuery, release: mockRelease });
+
+    const { initDatabase } = await import("../services/db.js");
+    await expect(initDatabase()).rejects.toThrow("DDL failed");
+    expect(mockRelease).toHaveBeenCalledWith();
+  });
+});
+
+describe("cleanupStaleData", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockConnect.mockReset();
+    mockQuery.mockReset();
+    mockRelease.mockReset();
+  });
+
+  it("deletes expired verifications and orphaned usage, returns counts", async () => {
+    // queryWithRetry: connect → query → release for each call
+    const client = { query: mockQuery, release: mockRelease };
+    mockConnect.mockResolvedValue(client);
+
+    // First queryWithRetry call: expired verifications
+    mockQuery.mockResolvedValueOnce({ rows: [{ email: "a@b.com" }, { email: "c@d.com" }], rowCount: 2 });
+    // Second queryWithRetry call: orphaned usage
+    mockQuery.mockResolvedValueOnce({ rows: [{ key: "old_key" }], rowCount: 1 });
+
+    const { cleanupStaleData } = await import("../services/db.js");
+    const result = await cleanupStaleData();
+
+    expect(result).toEqual({ expiredVerifications: 2, orphanedUsage: 1 });
+  });
+
+  it("returns zeros when nothing to clean", async () => {
+    const client = { query: mockQuery, release: mockRelease };
+    mockConnect.mockResolvedValue(client);
+
+    mockQuery.mockResolvedValue({ rows: [], rowCount: 0 });
+
+    const { cleanupStaleData } = await import("../services/db.js");
+    const result = await cleanupStaleData();
+
+    expect(result).toEqual({ expiredVerifications: 0, orphanedUsage: 0 });
+  });
+});
diff --git a/src/__tests__/keys-coverage.test.ts b/src/__tests__/keys-coverage.test.ts
new file mode 100644
index 0000000..ece0678
--- /dev/null
+++ b/src/__tests__/keys-coverage.test.ts
@@ -0,0 +1,129 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+
+// Unmock keys service — we want real implementation
+vi.unmock("../services/keys.js");
+
+// DB is still mocked by setup.ts
+import { queryWithRetry } from "../services/db.js";
+
+describe("keys service — uncovered branches", () => {
+  let keys: typeof import("../services/keys.js");
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    keys = await import("../services/keys.js");
+  });
+
+  describe("loadKeys error path", () => {
+    it("sets keysCache to empty array when query fails (ignoring env seeds)", async () => {
+      const origKeys = process.env.API_KEYS;
+      delete process.env.API_KEYS;
+
+      vi.mocked(queryWithRetry).mockRejectedValueOnce(new Error("DB down"));
+      await keys.loadKeys();
+
+      // Cache should be empty — no keys valid
+      expect(keys.isValidKey("anything")).toBe(false);
+      expect(keys.getAllKeys()).toEqual([]);
+
+      process.env.API_KEYS = origKeys;
+    });
+  });
+
+  describe("loadKeys with API_KEYS env var", () => {
+    const origEnv = process.env.API_KEYS;
+
+    afterEach(() => {
+      if (origEnv !== undefined) {
+        process.env.API_KEYS = origEnv;
+      } else {
+        delete process.env.API_KEYS;
+      }
+    });
+
+    it("seeds keys from API_KEYS env var into cache and upserts to DB", async () => {
+      process.env.API_KEYS = "seed_key_1,seed_key_2";
+
+      // loadKeys query returns empty (no existing keys)
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 0 } as any);
+      // Two upsert calls for seed keys
+      vi.mocked(queryWithRetry).mockResolvedValue({ rows: [], rowCount: 0 } as any);
+
+      await keys.loadKeys();
+
+      expect(keys.isValidKey("seed_key_1")).toBe(true);
+      expect(keys.isValidKey("seed_key_2")).toBe(true);
+      expect(keys.getKeyInfo("seed_key_1")?.tier).toBe("pro");
+      expect(keys.getKeyInfo("seed_key_1")?.email).toBe("seed@docfast.dev");
+
+      // 1 SELECT + 2 upserts
+      expect(vi.mocked(queryWithRetry)).toHaveBeenCalledTimes(3);
+    });
+
+    it("does not duplicate seed keys already in cache", async () => {
+      process.env.API_KEYS = "existing_key";
+
+      // loadKeys returns the key already in DB
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({
+        rows: [{ key: "existing_key", tier: "pro", email: "x@y.com", created_at: "2025-01-01T00:00:00Z", stripe_customer_id: null }],
+        rowCount: 1,
+      } as any);
+
+      await keys.loadKeys();
+
+      // Should not upsert — key already exists
+      expect(vi.mocked(queryWithRetry)).toHaveBeenCalledTimes(1);
+      expect(keys.getAllKeys()).toHaveLength(1);
+    });
+  });
+
+  describe("createFreeKey — email already exists in cache", () => {
+    it("returns existing free key for same email without DB insert", async () => {
+      // Load empty cache
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 0 } as any);
+      await keys.loadKeys();
+
+      // First create
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 1 } as any);
+      const first = await keys.createFreeKey("dup@test.com");
+
+      // Second call should return cached key without hitting DB
+      vi.mocked(queryWithRetry).mockClear();
+      const second = await keys.createFreeKey("dup@test.com");
+
+      expect(second.key).toBe(first.key);
+      expect(vi.mocked(queryWithRetry)).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("createProKey — customer already in cache", () => {
+    it("updates tier to pro and returns existing entry", async () => {
+      // Load cache with a free key that has a stripe customer
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({
+        rows: [{
+          key: "df_free_old",
+          tier: "free",
+          email: "user@test.com",
+          created_at: "2025-01-01T00:00:00Z",
+          stripe_customer_id: "cus_existing",
+        }],
+        rowCount: 1,
+      } as any);
+      await keys.loadKeys();
+
+      // Clear mock call history from loadKeys
+      vi.mocked(queryWithRetry).mockClear();
+      // The UPDATE query
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 1 } as any);
+
+      const result = await keys.createProKey("user@test.com", "cus_existing");
+
+      expect(result.key).toBe("df_free_old");
+      expect(result.tier).toBe("pro");
+      // Should have called UPDATE, not INSERT
+      const updateCall = vi.mocked(queryWithRetry).mock.calls[0];
+      expect(updateCall[0]).toContain("UPDATE");
+    });
+  });
+});