289 lines
10 KiB
Python
289 lines
10 KiB
Python
"""
|
|
Tests for screenshot module.
|
|
|
|
Tests async screenshot capture configuration and functionality.
|
|
"""
|
|
|
|
import pytest
|
|
from unittest.mock import Mock, AsyncMock, patch, MagicMock
|
|
import asyncio
|
|
import tempfile
|
|
import os
|
|
|
|
from neopig.screenshot import ScreenshotConfig, ScreenshotCapture
|
|
|
|
|
|
class TestScreenshotConfig:
|
|
"""Test ScreenshotConfig dataclass."""
|
|
|
|
def test_default_config(self):
|
|
"""Test default configuration values."""
|
|
config = ScreenshotConfig()
|
|
|
|
assert config.enabled is True
|
|
assert config.width == 1024
|
|
assert config.height == 768
|
|
assert config.delay == 1000
|
|
assert config.timeout == 30000
|
|
assert config.user_agent is None
|
|
assert config.engine is None
|
|
assert config.full_page is True
|
|
assert config.format == 'jpeg'
|
|
assert config.quality == 93
|
|
|
|
def test_custom_config(self):
|
|
"""Test custom configuration values."""
|
|
config = ScreenshotConfig(
|
|
enabled=True,
|
|
width=1920,
|
|
height=1080,
|
|
delay=2000,
|
|
user_agent="CustomBot/1.0",
|
|
engine="wkhtmltoimage",
|
|
)
|
|
|
|
assert config.enabled is True
|
|
assert config.width == 1920
|
|
assert config.height == 1080
|
|
assert config.delay == 2000
|
|
assert config.user_agent == "CustomBot/1.0"
|
|
assert config.engine == "wkhtmltoimage"
|
|
|
|
|
|
class TestScreenshotCapture:
|
|
"""Test ScreenshotCapture class."""
|
|
|
|
def test_capture_initialization_default(self):
|
|
"""Test ScreenshotCapture with default config."""
|
|
capture = ScreenshotCapture()
|
|
|
|
assert capture.config is not None
|
|
assert capture.config.enabled is True
|
|
assert capture._engine is None
|
|
assert capture._initialized is False
|
|
|
|
def test_capture_initialization_custom(self):
|
|
"""Test ScreenshotCapture with custom config."""
|
|
config = ScreenshotConfig(enabled=True, width=800, height=600)
|
|
capture = ScreenshotCapture(config=config)
|
|
|
|
assert capture.config.enabled is True
|
|
assert capture.config.width == 800
|
|
assert capture.config.height == 600
|
|
|
|
|
|
class TestScreenshotCaptureAvailability:
|
|
"""Test screenshot engine availability checking."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_available_caches_result(self):
|
|
"""Test that availability check caches initialization."""
|
|
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
|
|
|
# Mock the engine selection and creation
|
|
with patch.object(capture, '_get_available_engines', new_callable=AsyncMock) as mock_engines:
|
|
mock_engines.return_value = [{'name': 'wkhtmltoimage'}]
|
|
with patch.object(capture, '_select_engine', new_callable=AsyncMock) as mock_select:
|
|
mock_select.return_value = 'wkhtmltoimage'
|
|
with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread:
|
|
mock_engine = Mock()
|
|
mock_thread.return_value = mock_engine
|
|
|
|
# First call initializes
|
|
result1 = await capture.is_available()
|
|
# Second call uses cache
|
|
result2 = await capture.is_available()
|
|
|
|
assert result1 == result2
|
|
# Should only initialize once
|
|
assert mock_select.call_count == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_available_when_engine_found(self):
|
|
"""Test availability when an engine is found."""
|
|
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
|
|
|
with patch.object(capture, '_select_engine', new_callable=AsyncMock) as mock_select:
|
|
mock_select.return_value = 'wkhtmltoimage'
|
|
with patch('asyncio.to_thread', new_callable=AsyncMock) as mock_thread:
|
|
mock_engine = Mock()
|
|
mock_thread.return_value = mock_engine
|
|
|
|
result = await capture.is_available()
|
|
assert result is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_is_available_when_no_engine(self):
|
|
"""Test availability when no engine is available."""
|
|
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
|
|
|
with patch.object(capture, '_select_engine', new_callable=AsyncMock) as mock_select:
|
|
mock_select.return_value = None
|
|
|
|
result = await capture.is_available()
|
|
assert result is False
|
|
|
|
|
|
class TestScreenshotCaptureMethod:
|
|
"""Test the capture method."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_capture_disabled_returns_none(self):
|
|
"""Test capture returns None when disabled."""
|
|
config = ScreenshotConfig(enabled=False)
|
|
capture = ScreenshotCapture(config=config)
|
|
|
|
result = await capture.capture("https://example.com")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_capture_unavailable_returns_none(self):
|
|
"""Test capture returns None when no engine available."""
|
|
config = ScreenshotConfig(enabled=True)
|
|
capture = ScreenshotCapture(config=config)
|
|
|
|
with patch.object(capture, 'is_available', new_callable=AsyncMock, return_value=False):
|
|
result = await capture.capture("https://example.com")
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_capture_timeout(self):
|
|
"""Test capture handles timeout."""
|
|
config = ScreenshotConfig(enabled=True, timeout=1000)
|
|
capture = ScreenshotCapture(config=config)
|
|
capture._initialized = True
|
|
capture._engine = Mock()
|
|
capture._engine_name = 'test'
|
|
|
|
# Make engine.capture raise TimeoutError
|
|
async def slow_capture(*args):
|
|
raise asyncio.TimeoutError()
|
|
|
|
capture._engine.capture = slow_capture
|
|
|
|
result = await capture.capture("https://slow-site.com")
|
|
assert result is None
|
|
|
|
|
|
class TestScreenshotCaptureToFile:
|
|
"""Test capture_to_file method."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_capture_to_file_success(self):
|
|
"""Test successful capture to file."""
|
|
config = ScreenshotConfig(enabled=True)
|
|
capture = ScreenshotCapture(config=config)
|
|
|
|
fake_result = {
|
|
'data': b'fake png data',
|
|
'md5_hash': 'abc123',
|
|
'mime_type': 'image/png',
|
|
'size': 13,
|
|
}
|
|
|
|
with patch.object(capture, 'capture', new_callable=AsyncMock, return_value=fake_result):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
output_path = os.path.join(tmpdir, 'screenshot.png')
|
|
result = await capture.capture_to_file("https://example.com", output_path)
|
|
|
|
assert result is True
|
|
assert os.path.exists(output_path)
|
|
with open(output_path, 'rb') as f:
|
|
assert f.read() == b'fake png data'
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_capture_to_file_failure(self):
|
|
"""Test capture_to_file when capture fails."""
|
|
config = ScreenshotConfig(enabled=True)
|
|
capture = ScreenshotCapture(config=config)
|
|
|
|
with patch.object(capture, 'capture', new_callable=AsyncMock, return_value=None):
|
|
result = await capture.capture_to_file("https://example.com", "/tmp/test.png")
|
|
assert result is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_capture_to_file_write_error(self):
|
|
"""Test capture_to_file handles write errors."""
|
|
config = ScreenshotConfig(enabled=True)
|
|
capture = ScreenshotCapture(config=config)
|
|
|
|
fake_result = {
|
|
'data': b'fake png data',
|
|
'md5_hash': 'abc123',
|
|
'mime_type': 'image/png',
|
|
'size': 13,
|
|
}
|
|
|
|
with patch.object(capture, 'capture', new_callable=AsyncMock, return_value=fake_result):
|
|
# Try to write to invalid path
|
|
result = await capture.capture_to_file(
|
|
"https://example.com",
|
|
"/nonexistent/path/screenshot.png"
|
|
)
|
|
assert result is False
|
|
|
|
|
|
class TestScreenshotResultStructure:
|
|
"""Test screenshot result structure."""
|
|
|
|
def test_expected_result_keys(self):
|
|
"""Document expected keys in capture result."""
|
|
expected_keys = ['data', 'md5_hash', 'mime_type', 'size', 'source_uri', 'engine']
|
|
|
|
# Just document the structure - actual capture requires uri2png
|
|
result = {
|
|
'data': b'png bytes',
|
|
'md5_hash': 'd41d8cd98f00b204e9800998ecf8427e',
|
|
'mime_type': 'image/png',
|
|
'size': 1234,
|
|
'source_uri': 'https://example.com',
|
|
'engine': 'wkhtmltoimage',
|
|
}
|
|
|
|
for key in expected_keys:
|
|
assert key in result
|
|
|
|
|
|
class TestEngineSelection:
|
|
"""Test screenshot engine selection logic."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_select_engine_prefers_native(self):
|
|
"""Test that native engines are preferred over browser-based."""
|
|
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
|
|
|
with patch.object(capture, '_get_available_engines', new_callable=AsyncMock) as mock_engines:
|
|
# Both native and browser engines available
|
|
mock_engines.return_value = [
|
|
{'name': 'playwright-chromium'},
|
|
{'name': 'wkhtmltoimage'},
|
|
]
|
|
with patch('shutil.which', return_value='/usr/bin/wkhtmltoimage'):
|
|
engine = await capture._select_engine()
|
|
assert engine == 'wkhtmltoimage'
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_select_engine_uses_specified(self):
|
|
"""Test that specified engine is used if available."""
|
|
config = ScreenshotConfig(enabled=True, engine='playwright-webkit')
|
|
capture = ScreenshotCapture(config)
|
|
|
|
with patch.object(capture, '_get_available_engines', new_callable=AsyncMock) as mock_engines:
|
|
mock_engines.return_value = [
|
|
{'name': 'wkhtmltoimage'},
|
|
{'name': 'playwright-webkit'},
|
|
]
|
|
engine = await capture._select_engine()
|
|
assert engine == 'playwright-webkit'
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_engine_name(self):
|
|
"""Test get_engine_name returns active engine."""
|
|
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
|
capture._engine_name = 'cutycapt'
|
|
|
|
assert capture.get_engine_name() == 'cutycapt'
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pytest.main([__file__, '-v'])
|