test: add auth, rate-limit, and keys service tests

src/__tests__/auth.test.ts

@@ -0,0 +1,85 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { authMiddleware } from "../middleware/auth.js";
+import { isValidKey, getKeyInfo } from "../services/keys.js";
+
+const mockJson = vi.fn();
+const mockStatus = vi.fn(() => ({ json: mockJson }));
+const mockNext = vi.fn();
+
+function makeReq(headers: Record<string, string> = {}): any {
+  return { headers };
+}
+
+function makeRes(): any {
+  mockJson.mockClear();
+  mockStatus.mockClear();
+  return { status: mockStatus, json: mockJson };
+}
+
+describe("authMiddleware", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("returns 401 when no auth header and no x-api-key", () => {
+    const req = makeReq();
+    const res = makeRes();
+    authMiddleware(req, res, mockNext);
+    expect(mockStatus).toHaveBeenCalledWith(401);
+    expect(mockJson).toHaveBeenCalledWith(
+      expect.objectContaining({ error: expect.stringContaining("Missing API key") })
+    );
+    expect(mockNext).not.toHaveBeenCalled();
+  });
+
+  it("returns 403 when Bearer token is invalid", () => {
+    vi.mocked(isValidKey).mockReturnValueOnce(false);
+    const req = makeReq({ authorization: "Bearer bad-key" });
+    const res = makeRes();
+    authMiddleware(req, res, mockNext);
+    expect(mockStatus).toHaveBeenCalledWith(403);
+    expect(mockJson).toHaveBeenCalledWith({ error: "Invalid API key" });
+    expect(mockNext).not.toHaveBeenCalled();
+  });
+
+  it("returns 403 when x-api-key is invalid", () => {
+    vi.mocked(isValidKey).mockReturnValueOnce(false);
+    const req = makeReq({ "x-api-key": "bad-key" });
+    const res = makeRes();
+    authMiddleware(req, res, mockNext);
+    expect(mockStatus).toHaveBeenCalledWith(403);
+    expect(mockNext).not.toHaveBeenCalled();
+  });
+
+  it("calls next() and attaches apiKeyInfo when Bearer token is valid", () => {
+    const info = { key: "test-key", tier: "pro", email: "t@t.com", createdAt: "2025-01-01" };
+    vi.mocked(isValidKey).mockReturnValueOnce(true);
+    vi.mocked(getKeyInfo).mockReturnValueOnce(info as any);
+    const req = makeReq({ authorization: "Bearer test-key" });
+    const res = makeRes();
+    authMiddleware(req, res, mockNext);
+    expect(mockNext).toHaveBeenCalled();
+    expect((req as any).apiKeyInfo).toEqual(info);
+  });
+
+  it("calls next() and attaches apiKeyInfo when x-api-key is valid", () => {
+    const info = { key: "xkey", tier: "free", email: "x@t.com", createdAt: "2025-01-01" };
+    vi.mocked(isValidKey).mockReturnValueOnce(true);
+    vi.mocked(getKeyInfo).mockReturnValueOnce(info as any);
+    const req = makeReq({ "x-api-key": "xkey" });
+    const res = makeRes();
+    authMiddleware(req, res, mockNext);
+    expect(mockNext).toHaveBeenCalled();
+    expect((req as any).apiKeyInfo).toEqual(info);
+  });
+
+  it("prefers Authorization header over x-api-key when both present", () => {
+    vi.mocked(isValidKey).mockReturnValueOnce(true);
+    vi.mocked(getKeyInfo).mockReturnValueOnce({ key: "bearer-key" } as any);
+    const req = makeReq({ authorization: "Bearer bearer-key", "x-api-key": "header-key" });
+    const res = makeRes();
+    authMiddleware(req, res, mockNext);
+    expect(isValidKey).toHaveBeenCalledWith("bearer-key");
+    expect((req as any).apiKeyInfo).toEqual({ key: "bearer-key" });
+  });
+});

src/__tests__/keys.test.ts

@@ -0,0 +1,108 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+
+// Unmock keys service — we want to test the real implementation
+vi.unmock("../services/keys.js");
+
+// DB is still mocked by setup.ts
+import { queryWithRetry } from "../services/db.js";
+
+describe("keys service", () => {
+  let keys: typeof import("../services/keys.js");
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    // Re-import to get fresh cache
+    keys = await import("../services/keys.js");
+  });
+
+  describe("after loadKeys", () => {
+    const mockRows = [
+      { key: "df_free_abc", tier: "free", email: "a@b.com", created_at: "2025-01-01T00:00:00Z", stripe_customer_id: null },
+      { key: "df_pro_xyz", tier: "pro", email: "pro@b.com", created_at: "2025-01-01T00:00:00Z", stripe_customer_id: "cus_123" },
+    ];
+
+    beforeEach(async () => {
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: mockRows, rowCount: 2 } as any);
+      await keys.loadKeys();
+    });
+
+    it("isValidKey returns true for cached keys", () => {
+      expect(keys.isValidKey("df_free_abc")).toBe(true);
+      expect(keys.isValidKey("df_pro_xyz")).toBe(true);
+    });
+
+    it("isValidKey returns false for unknown keys", () => {
+      expect(keys.isValidKey("unknown")).toBe(false);
+    });
+
+    it("isProKey returns true for pro tier, false for free", () => {
+      expect(keys.isProKey("df_pro_xyz")).toBe(true);
+      expect(keys.isProKey("df_free_abc")).toBe(false);
+    });
+
+    it("getKeyInfo returns correct ApiKey object", () => {
+      const info = keys.getKeyInfo("df_pro_xyz");
+      expect(info).toEqual({
+        key: "df_pro_xyz",
+        tier: "pro",
+        email: "pro@b.com",
+        createdAt: "2025-01-01T00:00:00Z",
+        stripeCustomerId: "cus_123",
+      });
+    });
+
+    it("getKeyInfo returns undefined for unknown key", () => {
+      expect(keys.getKeyInfo("nope")).toBeUndefined();
+    });
+  });
+
+  describe("createFreeKey", () => {
+    beforeEach(async () => {
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 0 } as any);
+      await keys.loadKeys();
+    });
+
+    it("creates key with df_free prefix", async () => {
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 1 } as any);
+      const result = await keys.createFreeKey("new@test.com");
+      expect(result.key).toMatch(/^df_free_/);
+      expect(result.tier).toBe("free");
+      expect(result.email).toBe("new@test.com");
+    });
+
+    it("returns existing key for same email", async () => {
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 1 } as any);
+      const first = await keys.createFreeKey("dup@test.com");
+      const second = await keys.createFreeKey("dup@test.com");
+      expect(second.key).toBe(first.key);
+    });
+  });
+
+  describe("createProKey", () => {
+    beforeEach(async () => {
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 0 } as any);
+      await keys.loadKeys();
+    });
+
+    it("uses UPSERT and returns key", async () => {
+      const returnedRow = {
+        key: "df_pro_newkey",
+        tier: "pro",
+        email: "pro@test.com",
+        created_at: "2025-06-01T00:00:00Z",
+        stripe_customer_id: "cus_new",
+      };
+      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [returnedRow], rowCount: 1 } as any);
+      
+      const result = await keys.createProKey("pro@test.com", "cus_new");
+      expect(result.tier).toBe("pro");
+      expect(result.stripeCustomerId).toBe("cus_new");
+      
+      const call = vi.mocked(queryWithRetry).mock.calls.find(
+        (c) => typeof c[0] === "string" && c[0].includes("ON CONFLICT")
+      );
+      expect(call).toBeTruthy();
+    });
+  });
+});

src/__tests__/pdfRateLimit.test.ts

@@ -0,0 +1,153 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { isProKey } from "../services/keys.js";
+
+// We need to import the middleware fresh to reset internal state.
+// The global setup already mocks keys service.
+
+// Since the module has internal state (rateLimitStore, activePdfCount),
+// we need to be careful about test isolation.
+
+const mockNext = vi.fn();
+const headers: Record<string, string> = {};
+const mockSet = vi.fn((k: string, v: string) => { headers[k] = v; });
+const mockJson = vi.fn();
+const mockStatus = vi.fn(() => ({ json: mockJson }));
+
+function makeReq(key = "test-key", tier = "free"): any {
+  return {
+    apiKeyInfo: { key, tier, email: "t@t.com", createdAt: "2025-01-01" },
+    headers: {},
+  };
+}
+
+function makeRes(): any {
+  Object.keys(headers).forEach((k) => delete headers[k]);
+  mockSet.mockClear();
+  mockJson.mockClear();
+  mockStatus.mockClear();
+  return { set: mockSet, status: mockStatus, json: mockJson };
+}
+
+describe("pdfRateLimitMiddleware", () => {
+  // Re-import module each test to reset internal state
+  let pdfRateLimitMiddleware: any;
+
+  beforeEach(async () => {
+    vi.clearAllMocks();
+    // Reset module to clear internal rateLimitStore and counters
+    vi.resetModules();
+    const mod = await import("../middleware/pdfRateLimit.js");
+    pdfRateLimitMiddleware = mod.pdfRateLimitMiddleware;
+  });
+
+  it("sets rate limit headers on response", () => {
+    vi.mocked(isProKey).mockReturnValue(false);
+    const req = makeReq("key-a");
+    const res = makeRes();
+    pdfRateLimitMiddleware(req, res, mockNext);
+    expect(mockSet).toHaveBeenCalledWith("X-RateLimit-Limit", "10");
+    expect(mockSet).toHaveBeenCalledWith("X-RateLimit-Remaining", expect.any(String));
+    expect(mockSet).toHaveBeenCalledWith("X-RateLimit-Reset", expect.any(String));
+  });
+
+  it("allows requests under rate limit", () => {
+    vi.mocked(isProKey).mockReturnValue(false);
+    const req = makeReq("key-b");
+    const res = makeRes();
+    pdfRateLimitMiddleware(req, res, mockNext);
+    expect(mockNext).toHaveBeenCalled();
+    expect(mockStatus).not.toHaveBeenCalled();
+  });
+
+  it("returns 429 with Retry-After when free tier rate limit exceeded (10/min)", () => {
+    vi.mocked(isProKey).mockReturnValue(false);
+    // Exhaust 10 requests
+    for (let i = 0; i < 10; i++) {
+      const req = makeReq("key-c");
+      const res = makeRes();
+      pdfRateLimitMiddleware(req, res, mockNext);
+    }
+    // 11th should be rejected
+    mockNext.mockClear();
+    const req = makeReq("key-c");
+    const res = makeRes();
+    pdfRateLimitMiddleware(req, res, mockNext);
+    expect(mockStatus).toHaveBeenCalledWith(429);
+    expect(mockSet).toHaveBeenCalledWith("Retry-After", expect.any(String));
+    expect(mockNext).not.toHaveBeenCalled();
+  });
+
+  it("returns 429 for pro tier at 30/min limit", () => {
+    vi.mocked(isProKey).mockReturnValue(true);
+    for (let i = 0; i < 30; i++) {
+      const req = makeReq("key-d");
+      const res = makeRes();
+      pdfRateLimitMiddleware(req, res, mockNext);
+    }
+    mockNext.mockClear();
+    const req = makeReq("key-d");
+    const res = makeRes();
+    pdfRateLimitMiddleware(req, res, mockNext);
+    expect(mockStatus).toHaveBeenCalledWith(429);
+    expect(mockNext).not.toHaveBeenCalled();
+  });
+
+  it("resets rate limit after window expires", async () => {
+    vi.mocked(isProKey).mockReturnValue(false);
+    // Use fake timers
+    vi.useFakeTimers();
+    try {
+      for (let i = 0; i < 10; i++) {
+        pdfRateLimitMiddleware(makeReq("key-e"), makeRes(), mockNext);
+      }
+      // Should be blocked
+      mockNext.mockClear();
+      pdfRateLimitMiddleware(makeReq("key-e"), makeRes(), mockNext);
+      expect(mockNext).not.toHaveBeenCalled();
+
+      // Advance past window (60s)
+      vi.advanceTimersByTime(61_000);
+
+      mockNext.mockClear();
+      pdfRateLimitMiddleware(makeReq("key-e"), makeRes(), mockNext);
+      expect(mockNext).toHaveBeenCalled();
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("returns 429 QUEUE_FULL when concurrency queue is full", async () => {
+    vi.mocked(isProKey).mockReturnValue(false);
+    // Access getConcurrencyStats to verify
+    const mod = await import("../middleware/pdfRateLimit.js");
+    
+    // Fill up concurrent slots (3) and queue (10) by acquiring slots without releasing
+    // We need 3 active + 10 queued = 13 acquires without release
+    const req = makeReq("key-f");
+    const res = makeRes();
+    pdfRateLimitMiddleware(req, res, mockNext);
+    
+    // The middleware attaches acquirePdfSlot; fill slots
+    const promises: Promise<void>[] = [];
+    // Acquire 3 active slots
+    for (let i = 0; i < 3; i++) {
+      const r = makeReq(`fill-${i}`);
+      const s = makeRes();
+      pdfRateLimitMiddleware(r, s, vi.fn());
+      await (r as any).acquirePdfSlot();
+    }
+    // Fill queue with 10
+    for (let i = 0; i < 10; i++) {
+      const r = makeReq(`queue-${i}`);
+      const s = makeRes();
+      pdfRateLimitMiddleware(r, s, vi.fn());
+      promises.push((r as any).acquirePdfSlot());
+    }
+    
+    // Next acquire should throw QUEUE_FULL
+    const rFull = makeReq("key-full");
+    const sFull = makeRes();
+    pdfRateLimitMiddleware(rFull, sFull, vi.fn());
+    await expect((rFull as any).acquirePdfSlot()).rejects.toThrow("QUEUE_FULL");
+  });
+});