test: add unit tests for browser pool and screenshot services

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()
+      try {
-      const acquired = []
+        await mod.initBrowser()
-      for (let i = 0; i < 8; i++) {
+        const acquired = []
-        acquired.push(await mod.acquirePage())
+        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()
     })
   })
 

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<typeof createMockPage>
+  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()
+    })
+  })
+})