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')), emulateMediaFeatures: vi.fn().mockResolvedValue(undefined), addStyleTag: vi.fn().mockResolvedValue(undefined), } } 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('darkMode', () => { it('emulates prefers-color-scheme: dark when darkMode is true', async () => { await takeScreenshot({ url: 'https://example.com', darkMode: true }) expect(mockPage.emulateMediaFeatures).toHaveBeenCalledWith([ { name: 'prefers-color-scheme', value: 'dark' } ]) }) it('calls emulateMediaFeatures before goto', async () => { const callOrder: string[] = [] mockPage.emulateMediaFeatures.mockImplementation(async () => { callOrder.push('emulate') }) mockPage.goto.mockImplementation(async () => { callOrder.push('goto') }) await takeScreenshot({ url: 'https://example.com', darkMode: true }) expect(callOrder).toEqual(['emulate', 'goto']) }) it('does not emulate media features when darkMode is false', async () => { await takeScreenshot({ url: 'https://example.com', darkMode: false }) expect(mockPage.emulateMediaFeatures).not.toHaveBeenCalled() }) it('does not emulate media features when darkMode is not set', async () => { await takeScreenshot({ url: 'https://example.com' }) expect(mockPage.emulateMediaFeatures).not.toHaveBeenCalled() }) }) describe('hideSelectors', () => { it('injects style tag to hide selectors after page load', async () => { await takeScreenshot({ url: 'https://example.com', hideSelectors: ['.ad', '#banner'] }) expect(mockPage.addStyleTag).toHaveBeenCalledWith({ content: '.ad { display: none !important }\n#banner { display: none !important }' }) }) it('calls addStyleTag after goto', async () => { const callOrder: string[] = [] mockPage.goto.mockImplementation(async () => { callOrder.push('goto') }) mockPage.addStyleTag.mockImplementation(async () => { callOrder.push('addStyleTag') }) await takeScreenshot({ url: 'https://example.com', hideSelectors: ['.ad'] }) expect(callOrder.indexOf('goto')).toBeLessThan(callOrder.indexOf('addStyleTag')) }) it('does not inject style tag when hideSelectors is empty', async () => { await takeScreenshot({ url: 'https://example.com', hideSelectors: [] }) expect(mockPage.addStyleTag).not.toHaveBeenCalled() }) it('does not inject style tag when hideSelectors is not set', async () => { await takeScreenshot({ url: 'https://example.com' }) expect(mockPage.addStyleTag).not.toHaveBeenCalled() }) it('handles single selector', async () => { await takeScreenshot({ url: 'https://example.com', hideSelectors: ['.cookie-banner'] }) expect(mockPage.addStyleTag).toHaveBeenCalledWith({ content: '.cookie-banner { display: none !important }' }) }) }) describe('css parameter', () => { it('injects custom CSS via addStyleTag', async () => { await takeScreenshot({ url: 'https://example.com', css: 'body { background: red !important }' }) expect(mockPage.addStyleTag).toHaveBeenCalledWith({ content: 'body { background: red !important }' }) }) it('injects css after goto and waitForSelector', async () => { const callOrder: string[] = [] mockPage.goto.mockImplementation(async () => { callOrder.push('goto') }) mockPage.waitForSelector.mockImplementation(async () => { callOrder.push('waitForSelector') }) mockPage.addStyleTag.mockImplementation(async () => { callOrder.push('addStyleTag') }) await takeScreenshot({ url: 'https://example.com', css: 'body { color: blue }', waitForSelector: '#main' }) expect(callOrder.indexOf('goto')).toBeLessThan(callOrder.indexOf('addStyleTag')) expect(callOrder.indexOf('waitForSelector')).toBeLessThan(callOrder.indexOf('addStyleTag')) }) it('works alongside hideSelectors', async () => { await takeScreenshot({ url: 'https://example.com', css: 'body { color: blue }', hideSelectors: ['.ad'] }) // Both should result in addStyleTag calls expect(mockPage.addStyleTag).toHaveBeenCalledWith({ content: 'body { color: blue }' }) expect(mockPage.addStyleTag).toHaveBeenCalledWith({ content: '.ad { display: none !important }' }) }) it('works alongside darkMode', async () => { await takeScreenshot({ url: 'https://example.com', css: 'h1 { font-size: 48px }', darkMode: true }) expect(mockPage.emulateMediaFeatures).toHaveBeenCalledWith([{ name: 'prefers-color-scheme', value: 'dark' }]) expect(mockPage.addStyleTag).toHaveBeenCalledWith({ content: 'h1 { font-size: 48px }' }) }) it('does not inject style tag when css is not set', async () => { await takeScreenshot({ url: 'https://example.com' }) expect(mockPage.addStyleTag).not.toHaveBeenCalled() }) it('does not inject style tag when css is empty string', async () => { await takeScreenshot({ url: 'https://example.com', css: '' }) expect(mockPage.addStyleTag).not.toHaveBeenCalled() }) }) 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() }) }) })