SnapAPI/src/services/__tests__/screenshot.test.ts

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<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('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()
    })
  })
})