docfast/src/__tests__/pdfRateLimit.test.ts

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");
  });
});