test: add browser pool unit tests

src/__tests__/browser-pool.test.ts

@@ -0,0 +1,347 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+
+// Don't use the global mock — we test the real browser service
+vi.unmock("../services/browser.js");
+
+// Mock logger
+vi.mock("../services/logger.js", () => ({
+  default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
+}));
+
+function createMockPage() {
+  const page: any = {
+    setJavaScriptEnabled: vi.fn().mockResolvedValue(undefined),
+    setContent: vi.fn().mockResolvedValue(undefined),
+    addStyleTag: vi.fn().mockResolvedValue(undefined),
+    pdf: vi.fn().mockResolvedValue(Buffer.from("%PDF-1.4 test")),
+    goto: vi.fn().mockResolvedValue(undefined),
+    close: vi.fn().mockResolvedValue(undefined),
+    setRequestInterception: vi.fn().mockResolvedValue(undefined),
+    removeAllListeners: vi.fn().mockReturnThis(),
+    createCDPSession: vi.fn().mockResolvedValue({
+      send: vi.fn().mockResolvedValue(undefined),
+      detach: vi.fn().mockResolvedValue(undefined),
+    }),
+    cookies: vi.fn().mockResolvedValue([]),
+    deleteCookie: vi.fn(),
+    on: vi.fn(),
+    newPage: vi.fn(),
+  };
+  return page;
+}
+
+function createMockBrowser(pagesPerBrowser = 8) {
+  const pages = Array.from({ length: pagesPerBrowser }, () => createMockPage());
+  let pageIndex = 0;
+  const browser: any = {
+    newPage: vi.fn().mockImplementation(() => Promise.resolve(pages[pageIndex++] || createMockPage())),
+    close: vi.fn().mockResolvedValue(undefined),
+    _pages: pages,
+  };
+  return browser;
+}
+
+// We need to set env vars before importing
+process.env.BROWSER_COUNT = "2";
+process.env.PAGES_PER_BROWSER = "2"; // small for testing
+
+let mockBrowsers: any[] = [];
+let launchCallCount = 0;
+
+vi.mock("puppeteer", () => ({
+  default: {
+    launch: vi.fn().mockImplementation(() => {
+      const b = createMockBrowser(2);
+      mockBrowsers.push(b);
+      launchCallCount++;
+      return Promise.resolve(b);
+    }),
+  },
+}));
+
+describe("browser pool", () => {
+  let browserModule: typeof import("../services/browser.js");
+
+  beforeEach(async () => {
+    mockBrowsers = [];
+    launchCallCount = 0;
+    // Fresh import each test to reset module state (instances array)
+    vi.resetModules();
+    // Re-apply mocks after resetModules
+    vi.doMock("puppeteer", () => ({
+      default: {
+        launch: vi.fn().mockImplementation(() => {
+          const b = createMockBrowser(2);
+          mockBrowsers.push(b);
+          launchCallCount++;
+          return Promise.resolve(b);
+        }),
+      },
+    }));
+    vi.doMock("../services/logger.js", () => ({
+      default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
+    }));
+    browserModule = await import("../services/browser.js");
+  });
+
+  afterEach(async () => {
+    try {
+      await browserModule.closeBrowser();
+    } catch {}
+  });
+
+  describe("initBrowser / closeBrowser", () => {
+    it("launches BROWSER_COUNT browser instances", async () => {
+      await browserModule.initBrowser();
+      expect(launchCallCount).toBe(2);
+      expect(mockBrowsers).toHaveLength(2);
+    });
+
+    it("creates PAGES_PER_BROWSER pages per browser", async () => {
+      await browserModule.initBrowser();
+      for (const b of mockBrowsers) {
+        expect(b.newPage).toHaveBeenCalledTimes(2);
+      }
+    });
+
+    it("closeBrowser closes all pages and browsers", async () => {
+      await browserModule.initBrowser();
+      const allPages = mockBrowsers.flatMap((b: any) => b._pages.slice(0, 2));
+      await browserModule.closeBrowser();
+      for (const page of allPages) {
+        expect(page.close).toHaveBeenCalled();
+      }
+      for (const b of mockBrowsers) {
+        expect(b.close).toHaveBeenCalled();
+      }
+    });
+  });
+
+  describe("getPoolStats", () => {
+    it("returns correct structure after init", async () => {
+      await browserModule.initBrowser();
+      const stats = browserModule.getPoolStats();
+      expect(stats).toMatchObject({
+        poolSize: 4, // 2 browsers × 2 pages
+        totalPages: 4,
+        availablePages: 4,
+        queueDepth: 0,
+        pdfCount: 0,
+        restarting: false,
+      });
+      expect(stats.browsers).toHaveLength(2);
+      expect(stats.browsers[0]).toMatchObject({
+        id: 0,
+        available: 2,
+        pdfCount: 0,
+        restarting: false,
+      });
+    });
+
+    it("returns empty stats before init", () => {
+      const stats = browserModule.getPoolStats();
+      expect(stats.poolSize).toBe(0);
+      expect(stats.availablePages).toBe(0);
+      expect(stats.browsers).toHaveLength(0);
+    });
+  });
+
+  describe("renderPdf", () => {
+    it("generates a PDF buffer from HTML", async () => {
+      await browserModule.initBrowser();
+      const result = await browserModule.renderPdf("<h1>Hello</h1>");
+      expect(Buffer.isBuffer(result)).toBe(true);
+      expect(result.toString()).toContain("%PDF");
+    });
+
+    it("sets content and disables JS on the page", async () => {
+      await browserModule.initBrowser();
+      await browserModule.renderPdf("<h1>Test</h1>");
+      // Find a page that was used
+      const usedPage = mockBrowsers
+        .flatMap((b: any) => b._pages.slice(0, 2))
+        .find((p: any) => p.setContent.mock.calls.length > 0);
+      expect(usedPage).toBeDefined();
+      expect(usedPage.setJavaScriptEnabled).toHaveBeenCalledWith(false);
+      expect(usedPage.setContent).toHaveBeenCalledWith("<h1>Test</h1>", expect.objectContaining({ waitUntil: "domcontentloaded" }));
+      expect(usedPage.pdf).toHaveBeenCalled();
+    });
+
+    it("releases the page back to the pool after rendering", async () => {
+      await browserModule.initBrowser();
+      const statsBefore = browserModule.getPoolStats();
+      await browserModule.renderPdf("<p>test</p>");
+      // After render + recycle, page should be available again (async recycle)
+      // pdfCount should have incremented
+      const statsAfter = browserModule.getPoolStats();
+      expect(statsAfter.pdfCount).toBe(1);
+    });
+
+    it("passes options correctly to page.pdf()", async () => {
+      await browserModule.initBrowser();
+      await browserModule.renderPdf("<p>test</p>", {
+        format: "Letter",
+        landscape: true,
+        scale: 0.8,
+        margin: { top: "10mm", bottom: "10mm", left: "5mm", right: "5mm" },
+        displayHeaderFooter: true,
+        headerTemplate: "<div>Header</div>",
+        footerTemplate: "<div>Footer</div>",
+      });
+      const usedPage = mockBrowsers
+        .flatMap((b: any) => b._pages.slice(0, 2))
+        .find((p: any) => p.pdf.mock.calls.length > 0);
+      const pdfArgs = usedPage.pdf.mock.calls[0][0];
+      expect(pdfArgs.format).toBe("Letter");
+      expect(pdfArgs.landscape).toBe(true);
+      expect(pdfArgs.scale).toBe(0.8);
+      expect(pdfArgs.margin).toEqual({ top: "10mm", bottom: "10mm", left: "5mm", right: "5mm" });
+      expect(pdfArgs.displayHeaderFooter).toBe(true);
+      expect(pdfArgs.headerTemplate).toBe("<div>Header</div>");
+    });
+
+    it("still releases page if setContent throws (no pool leak)", async () => {
+      await browserModule.initBrowser();
+      // Make ALL pages' setContent throw so whichever is picked will fail
+      for (const b of mockBrowsers) {
+        for (const p of b._pages) {
+          p.setContent.mockRejectedValueOnce(new Error("render fail"));
+        }
+      }
+
+      await expect(browserModule.renderPdf("<bad>")).rejects.toThrow("render fail");
+      // pdfCount should still increment (releasePage was called in finally)
+      const stats = browserModule.getPoolStats();
+      expect(stats.pdfCount).toBe(1);
+    });
+
+    it("rejects with PDF_TIMEOUT after 30s", async () => {
+      await browserModule.initBrowser();
+      // Make ALL pages' setContent hang so whichever is picked will timeout
+      for (const b of mockBrowsers) {
+        for (const p of b._pages) {
+          p.setContent.mockImplementation(() => new Promise(() => {}));
+        }
+      }
+
+      vi.useFakeTimers();
+      const renderPromise = browserModule.renderPdf("<h1>slow</h1>");
+      await vi.advanceTimersByTimeAsync(30_001);
+
+      await expect(renderPromise).rejects.toThrow("PDF_TIMEOUT");
+      vi.useRealTimers();
+    });
+  });
+
+  describe("renderUrlPdf", () => {
+    it("navigates to URL and generates PDF", async () => {
+      await browserModule.initBrowser();
+      const result = await browserModule.renderUrlPdf("https://example.com");
+      expect(Buffer.isBuffer(result)).toBe(true);
+      const usedPage = mockBrowsers
+        .flatMap((b: any) => b._pages.slice(0, 2))
+        .find((p: any) => p.goto.mock.calls.length > 0);
+      expect(usedPage.goto).toHaveBeenCalledWith("https://example.com", expect.objectContaining({ waitUntil: "domcontentloaded" }));
+    });
+
+    it("sets up request interception for SSRF protection with hostResolverRules", async () => {
+      await browserModule.initBrowser();
+      await browserModule.renderUrlPdf("https://example.com", {
+        hostResolverRules: "MAP example.com 93.184.216.34",
+      });
+      const usedPage = mockBrowsers
+        .flatMap((b: any) => b._pages.slice(0, 2))
+        .find((p: any) => p.setRequestInterception.mock.calls.length > 0);
+      expect(usedPage).toBeDefined();
+      expect(usedPage.setRequestInterception).toHaveBeenCalledWith(true);
+      expect(usedPage.on).toHaveBeenCalledWith("request", expect.any(Function));
+    });
+
+    it("blocks requests to non-target hosts via request interception", async () => {
+      await browserModule.initBrowser();
+      await browserModule.renderUrlPdf("https://example.com", {
+        hostResolverRules: "MAP example.com 93.184.216.34",
+      });
+      const usedPage = mockBrowsers
+        .flatMap((b: any) => b._pages.slice(0, 2))
+        .find((p: any) => p.on.mock.calls.length > 0);
+
+      // Get the request handler
+      const requestHandler = usedPage.on.mock.calls.find((c: any) => c[0] === "request")[1];
+
+      // Simulate a request to a different host
+      const evilRequest = {
+        url: () => "http://169.254.169.254/metadata",
+        headers: () => ({}),
+        abort: vi.fn(),
+        continue: vi.fn(),
+      };
+      requestHandler(evilRequest);
+      expect(evilRequest.abort).toHaveBeenCalledWith("blockedbyclient");
+      expect(evilRequest.continue).not.toHaveBeenCalled();
+
+      // Simulate a request to the target host (HTTP - should rewrite)
+      const goodRequest = {
+        url: () => "http://example.com/page",
+        headers: () => ({ "accept": "text/html" }),
+        abort: vi.fn(),
+        continue: vi.fn(),
+      };
+      requestHandler(goodRequest);
+      expect(goodRequest.continue).toHaveBeenCalledWith(expect.objectContaining({
+        url: expect.stringContaining("93.184.216.34"),
+        headers: expect.objectContaining({ host: "example.com" }),
+      }));
+      expect(goodRequest.abort).not.toHaveBeenCalled();
+    });
+  });
+
+  describe("acquirePage queue", () => {
+    it("queues requests when all pages are busy and resolves when released", async () => {
+      await browserModule.initBrowser();
+      // Use all 4 pages
+      const p1 = browserModule.renderPdf("<p>1</p>");
+      const p2 = browserModule.renderPdf("<p>2</p>");
+      const p3 = browserModule.renderPdf("<p>3</p>");
+      const p4 = browserModule.renderPdf("<p>4</p>");
+
+      // Stats should show queue or reduced availability
+      // The 5th request should queue
+      // But since our mock pages resolve instantly, the first 4 may already be done
+      // Let's make pages hang to truly test queuing
+      await Promise.all([p1, p2, p3, p4]);
+
+      // Verify all rendered successfully
+      const stats = browserModule.getPoolStats();
+      expect(stats.pdfCount).toBe(4);
+    });
+
+    it("rejects with QUEUE_FULL after 30s timeout when all pages busy", async () => {
+      await browserModule.initBrowser();
+
+      // Make all pages hang
+      for (const b of mockBrowsers) {
+        for (const p of b._pages) {
+          p.setContent.mockImplementation(() => new Promise(() => {}));
+        }
+      }
+
+      // Consume all 4 pages (these will hang)
+      browserModule.renderPdf("<p>1</p>");
+      browserModule.renderPdf("<p>2</p>");
+      browserModule.renderPdf("<p>3</p>");
+      browserModule.renderPdf("<p>4</p>");
+
+      vi.useFakeTimers();
+      // 5th request should queue
+      const queued = browserModule.renderPdf("<p>5</p>");
+
+      // Advance past queue timeout
+      await vi.advanceTimersByTimeAsync(30_001);
+
+      await expect(queued).rejects.toThrow("QUEUE_FULL");
+
+      vi.useRealTimers();
+    });
+  });
+});