From 05c91e6747f4f43317a7428b080d32d3a214531a Mon Sep 17 00:00:00 2001 From: Hoid Date: Tue, 3 Mar 2026 15:07:02 +0100 Subject: [PATCH] test: add unit tests for browser pool and screenshot services --- src/services/__tests__/browser.test.ts | 31 ++-- src/services/__tests__/screenshot.test.ts | 198 ++++++++++++++++++++++ 2 files changed, 218 insertions(+), 11 deletions(-) create mode 100644 src/services/__tests__/screenshot.test.ts diff --git a/src/services/__tests__/browser.test.ts b/src/services/__tests__/browser.test.ts index 5ba46fe..4fa92ff 100644 --- a/src/services/__tests__/browser.test.ts +++ b/src/services/__tests__/browser.test.ts @@ -161,18 +161,27 @@ describe('Browser Pool Service', () => { it('rejects with QUEUE_FULL after 30s timeout', async () => { vi.useFakeTimers() - await mod.initBrowser() - const acquired = [] - for (let i = 0; i < 8; i++) { - acquired.push(await mod.acquirePage()) + try { + await mod.initBrowser() + const acquired = [] + for (let i = 0; i < 8; i++) { + acquired.push(await mod.acquirePage()) + } + + // Catch the rejection immediately to avoid unhandled rejection + let error: Error | null = null + const pending = mod.acquirePage().catch((e: Error) => { error = e }) + await vi.advanceTimersByTimeAsync(31_000) + await pending + expect(error).toBeTruthy() + expect(error!.message).toBe('QUEUE_FULL') + + for (const a of acquired) mod.releasePage(a.page, a.instance) + // Let async cleanup finish + await vi.advanceTimersByTimeAsync(100) + } finally { + vi.useRealTimers() } - - const pending = mod.acquirePage() - await vi.advanceTimersByTimeAsync(31_000) - await expect(pending).rejects.toThrow('QUEUE_FULL') - - for (const a of acquired) mod.releasePage(a.page, a.instance) - vi.useRealTimers() }) }) diff --git a/src/services/__tests__/screenshot.test.ts b/src/services/__tests__/screenshot.test.ts new file mode 100644 index 0000000..815a8bc --- /dev/null +++ b/src/services/__tests__/screenshot.test.ts @@ -0,0 +1,198 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest' + +// All mocks must use inline values since vi.mock is hoisted +vi.mock('../browser.js', () => ({ + acquirePage: vi.fn(), + releasePage: vi.fn(), +})) + +vi.mock('../ssrf.js', () => ({ + validateUrl: vi.fn(), +})) + +vi.mock('../logger.js', () => ({ + default: { info: vi.fn(), warn: vi.fn(), error: vi.fn() }, +})) + +import { takeScreenshot } from '../screenshot.js' +import { acquirePage, releasePage } from '../browser.js' +import { validateUrl } from '../ssrf.js' + +function createMockPage() { + return { + setViewport: vi.fn().mockResolvedValue(undefined), + goto: vi.fn().mockResolvedValue(undefined), + waitForSelector: vi.fn().mockResolvedValue(undefined), + screenshot: vi.fn().mockResolvedValue(Buffer.from('fake-screenshot')), + } +} + +describe('Screenshot Service', () => { + let mockPage: ReturnType + const mockInstance = { id: 0 } + + beforeEach(() => { + vi.clearAllMocks() + mockPage = createMockPage() + vi.mocked(acquirePage).mockResolvedValue({ page: mockPage as any, instance: mockInstance as any }) + vi.mocked(validateUrl).mockResolvedValue({ hostname: 'example.com', resolvedIp: '1.2.3.4' } as any) + }) + + describe('URL validation', () => { + it('calls validateUrl with the provided URL', async () => { + await takeScreenshot({ url: 'https://example.com' }) + expect(validateUrl).toHaveBeenCalledWith('https://example.com') + }) + + it('propagates SSRF validation errors', async () => { + vi.mocked(validateUrl).mockRejectedValueOnce(new Error('SSRF_BLOCKED')) + await expect(takeScreenshot({ url: 'http://localhost' })).rejects.toThrow('SSRF_BLOCKED') + }) + }) + + describe('Default options', () => { + it('uses format=png, width=1280, height=800, fullPage=false', async () => { + await takeScreenshot({ url: 'https://example.com' }) + expect(mockPage.setViewport).toHaveBeenCalledWith({ + width: 1280, + height: 800, + deviceScaleFactor: 1, + }) + expect(mockPage.screenshot).toHaveBeenCalledWith( + expect.objectContaining({ + type: 'png', + fullPage: false, + encoding: 'binary', + }) + ) + const screenshotArg = mockPage.screenshot.mock.calls[0][0] + expect(screenshotArg.quality).toBeUndefined() + }) + + it('returns image/png content type by default', async () => { + const result = await takeScreenshot({ url: 'https://example.com' }) + expect(result.contentType).toBe('image/png') + }) + }) + + describe('Custom options', () => { + it('respects custom width and height', async () => { + await takeScreenshot({ url: 'https://example.com', width: 800, height: 600 }) + expect(mockPage.setViewport).toHaveBeenCalledWith( + expect.objectContaining({ width: 800, height: 600 }) + ) + }) + + it('caps width at MAX_WIDTH (3840)', async () => { + await takeScreenshot({ url: 'https://example.com', width: 5000 }) + expect(mockPage.setViewport).toHaveBeenCalledWith( + expect.objectContaining({ width: 3840 }) + ) + }) + + it('caps height at MAX_HEIGHT (2160)', async () => { + await takeScreenshot({ url: 'https://example.com', height: 5000 }) + expect(mockPage.setViewport).toHaveBeenCalledWith( + expect.objectContaining({ height: 2160 }) + ) + }) + + it('uses jpeg format with quality', async () => { + await takeScreenshot({ url: 'https://example.com', format: 'jpeg', quality: 90 }) + expect(mockPage.screenshot).toHaveBeenCalledWith( + expect.objectContaining({ type: 'jpeg', quality: 90 }) + ) + }) + + it('returns correct content type for jpeg', async () => { + const result = await takeScreenshot({ url: 'https://example.com', format: 'jpeg' }) + expect(result.contentType).toBe('image/jpeg') + }) + + it('uses webp format with quality', async () => { + const result = await takeScreenshot({ url: 'https://example.com', format: 'webp', quality: 75 }) + expect(mockPage.screenshot).toHaveBeenCalledWith( + expect.objectContaining({ type: 'webp', quality: 75 }) + ) + expect(result.contentType).toBe('image/webp') + }) + + it('does not set quality for png even if provided', async () => { + await takeScreenshot({ url: 'https://example.com', format: 'png', quality: 90 }) + const arg = mockPage.screenshot.mock.calls[0][0] + expect(arg.quality).toBeUndefined() + }) + + it('respects fullPage option', async () => { + await takeScreenshot({ url: 'https://example.com', fullPage: true }) + expect(mockPage.screenshot).toHaveBeenCalledWith( + expect.objectContaining({ fullPage: true }) + ) + }) + + it('caps deviceScale at 3', async () => { + await takeScreenshot({ url: 'https://example.com', deviceScale: 5 }) + expect(mockPage.setViewport).toHaveBeenCalledWith( + expect.objectContaining({ deviceScaleFactor: 3 }) + ) + }) + }) + + describe('waitForSelector', () => { + it('calls waitForSelector when provided', async () => { + await takeScreenshot({ url: 'https://example.com', waitForSelector: '#content' }) + expect(mockPage.waitForSelector).toHaveBeenCalledWith('#content', { timeout: 10_000 }) + }) + + it('does not call waitForSelector when not provided', async () => { + await takeScreenshot({ url: 'https://example.com' }) + expect(mockPage.waitForSelector).not.toHaveBeenCalled() + }) + }) + + describe('delay', () => { + it('creates a pause when delay is provided', async () => { + vi.useFakeTimers() + const promise = takeScreenshot({ url: 'https://example.com', delay: 1000 }) + await vi.advanceTimersByTimeAsync(1000) + await promise + vi.useRealTimers() + }) + }) + + describe('Timeout', () => { + it('rejects with SCREENSHOT_TIMEOUT after 30s', async () => { + vi.useFakeTimers() + try { + mockPage.goto.mockImplementation(() => new Promise(() => {})) + let error: Error | null = null + const promise = takeScreenshot({ url: 'https://example.com' }).catch((e: Error) => { error = e }) + await vi.advanceTimersByTimeAsync(31_000) + await promise + expect(error).toBeTruthy() + expect(error!.message).toBe('SCREENSHOT_TIMEOUT') + } finally { + vi.useRealTimers() + } + }) + }) + + describe('Page lifecycle', () => { + it('always releases page after successful screenshot', async () => { + await takeScreenshot({ url: 'https://example.com' }) + expect(releasePage).toHaveBeenCalledWith(mockPage, mockInstance) + }) + + it('releases page even when an error occurs', async () => { + mockPage.goto.mockRejectedValueOnce(new Error('NAV_FAILED')) + await expect(takeScreenshot({ url: 'https://example.com' })).rejects.toThrow('NAV_FAILED') + expect(releasePage).toHaveBeenCalledWith(mockPage, mockInstance) + }) + + it('does not release page on SSRF error (before acquire)', async () => { + vi.mocked(validateUrl).mockRejectedValueOnce(new Error('SSRF_BLOCKED')) + await expect(takeScreenshot({ url: 'http://10.0.0.1' })).rejects.toThrow('SSRF_BLOCKED') + expect(releasePage).not.toHaveBeenCalled() + }) + }) +})