From 0283e9dae8bc136b309decac0319f26ee8cd3167 Mon Sep 17 00:00:00 2001 From: OpenClaw Date: Fri, 6 Mar 2026 08:05:45 +0100 Subject: [PATCH] test: add browser pool unit tests --- src/__tests__/browser-pool.test.ts | 347 +++++++++++++++++++++++++++++ 1 file changed, 347 insertions(+) create mode 100644 src/__tests__/browser-pool.test.ts diff --git a/src/__tests__/browser-pool.test.ts b/src/__tests__/browser-pool.test.ts new file mode 100644 index 0000000..1252f4c --- /dev/null +++ b/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("

Hello

"); + 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("

Test

"); + // 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("

Test

", 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("

test

"); + // 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("

test

", { + format: "Letter", + landscape: true, + scale: 0.8, + margin: { top: "10mm", bottom: "10mm", left: "5mm", right: "5mm" }, + displayHeaderFooter: true, + headerTemplate: "
Header
", + footerTemplate: "
Footer
", + }); + 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("
Header
"); + }); + + 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("")).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("

slow

"); + 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("

1

"); + const p2 = browserModule.renderPdf("

2

"); + const p3 = browserModule.renderPdf("

3

"); + const p4 = browserModule.renderPdf("

4

"); + + // 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("

1

"); + browserModule.renderPdf("

2

"); + browserModule.renderPdf("

3

"); + browserModule.renderPdf("

4

"); + + vi.useFakeTimers(); + // 5th request should queue + const queued = browserModule.renderPdf("

5

"); + + // Advance past queue timeout + await vi.advanceTimersByTimeAsync(30_001); + + await expect(queued).rejects.toThrow("QUEUE_FULL"); + + vi.useRealTimers(); + }); + }); +});