modified: archive.py
modified: tests/unit/test_screenshot.py
This commit is contained in:
parent
21e344f675
commit
36edc3b8dc
2 changed files with 117 additions and 111 deletions
30
archive.py
30
archive.py
|
|
@ -536,13 +536,16 @@ def list_files(subdir: str, pattern: str = '*') -> list:
|
|||
|
||||
|
||||
def search_pages(query: str, limit: int = 50) -> list:
|
||||
"""Search pages using FTS5."""
|
||||
"""Search pages using FTS5 with LIKE fallback."""
|
||||
if not DB_PATH or not DB_PATH.exists():
|
||||
return []
|
||||
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
c = conn.cursor()
|
||||
results = []
|
||||
try:
|
||||
# Try FTS5 with prefix matching (add * for partial matches)
|
||||
fts_query = ' '.join(f'"{word}"*' for word in query.split())
|
||||
c.execute("""
|
||||
SELECT p.path, p.title, snippet(pages_fts, 1, '<mark>', '</mark>', '...', 40)
|
||||
FROM pages_fts
|
||||
|
|
@ -550,12 +553,28 @@ def search_pages(query: str, limit: int = 50) -> list:
|
|||
WHERE pages_fts MATCH ?
|
||||
ORDER BY rank
|
||||
LIMIT ?
|
||||
""", (query, limit))
|
||||
""", (fts_query, limit))
|
||||
results = [{'path': r[0], 'title': r[1], 'snippet': r[2]} for r in c.fetchall()]
|
||||
except sqlite3.OperationalError:
|
||||
results = []
|
||||
finally:
|
||||
conn.close()
|
||||
pass
|
||||
|
||||
# Fallback to LIKE for substring matching if FTS5 found nothing
|
||||
if not results:
|
||||
try:
|
||||
like_q = f'%{query}%'
|
||||
c.execute("""
|
||||
SELECT path, title, substr(content, 1, 200) as snippet
|
||||
FROM pages
|
||||
WHERE title LIKE ? COLLATE NOCASE
|
||||
OR content LIKE ? COLLATE NOCASE
|
||||
OR path LIKE ? COLLATE NOCASE
|
||||
LIMIT ?
|
||||
""", (like_q, like_q, like_q, limit))
|
||||
results = [{'path': r[0], 'title': r[1], 'snippet': r[2] + '...'} for r in c.fetchall()]
|
||||
except sqlite3.OperationalError:
|
||||
pass
|
||||
|
||||
conn.close()
|
||||
return results
|
||||
|
||||
|
||||
|
|
@ -986,6 +1005,7 @@ async def main():
|
|||
width=args.screenshot_width,
|
||||
height=args.screenshot_height,
|
||||
engine=args.screenshot_engine,
|
||||
full_page=True, # Archive captures full page by default
|
||||
)
|
||||
|
||||
archiver = SiteArchiver(
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ class TestScreenshotConfig:
|
|||
assert config.height == 1024
|
||||
assert config.delay == 1000
|
||||
assert config.user_agent is None
|
||||
assert config.engine is None
|
||||
|
||||
def test_custom_config(self):
|
||||
"""Test custom configuration values."""
|
||||
|
|
@ -34,6 +35,7 @@ class TestScreenshotConfig:
|
|||
height=1080,
|
||||
delay=2000,
|
||||
user_agent="CustomBot/1.0",
|
||||
engine="wkhtmltoimage",
|
||||
)
|
||||
|
||||
assert config.enabled is True
|
||||
|
|
@ -41,6 +43,7 @@ class TestScreenshotConfig:
|
|||
assert config.height == 1080
|
||||
assert config.delay == 2000
|
||||
assert config.user_agent == "CustomBot/1.0"
|
||||
assert config.engine == "wkhtmltoimage"
|
||||
|
||||
|
||||
class TestScreenshotCapture:
|
||||
|
|
@ -52,7 +55,8 @@ class TestScreenshotCapture:
|
|||
|
||||
assert capture.config is not None
|
||||
assert capture.config.enabled is False
|
||||
assert capture._uri2png_available is None
|
||||
assert capture._engine is None
|
||||
assert capture._initialized is False
|
||||
|
||||
def test_capture_initialization_custom(self):
|
||||
"""Test ScreenshotCapture with custom config."""
|
||||
|
|
@ -65,63 +69,52 @@ class TestScreenshotCapture:
|
|||
|
||||
|
||||
class TestScreenshotCaptureAvailability:
|
||||
"""Test uri2png availability checking."""
|
||||
"""Test screenshot engine availability checking."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_caches_result(self):
|
||||
"""Test that availability check is cached."""
|
||||
capture = ScreenshotCapture()
|
||||
"""Test that availability check caches initialization."""
|
||||
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
||||
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.returncode = 0
|
||||
mock_proc.wait = AsyncMock(return_value=0)
|
||||
mock_exec.return_value = mock_proc
|
||||
# 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
|
||||
result1 = await capture.is_available()
|
||||
# Second call should use cache
|
||||
result2 = await capture.is_available()
|
||||
# First call initializes
|
||||
result1 = await capture.is_available()
|
||||
# Second call uses cache
|
||||
result2 = await capture.is_available()
|
||||
|
||||
assert result1 == result2
|
||||
# Should only call subprocess once
|
||||
assert mock_exec.call_count == 1
|
||||
assert result1 == result2
|
||||
# Should only initialize once
|
||||
assert mock_select.call_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_when_installed(self):
|
||||
"""Test availability when uri2png is installed."""
|
||||
capture = ScreenshotCapture()
|
||||
async def test_is_available_when_engine_found(self):
|
||||
"""Test availability when an engine is found."""
|
||||
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
||||
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.returncode = 0
|
||||
mock_proc.wait = AsyncMock(return_value=0)
|
||||
mock_exec.return_value = mock_proc
|
||||
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
|
||||
result = await capture.is_available()
|
||||
assert result is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_when_not_installed(self):
|
||||
"""Test availability when uri2png is not installed."""
|
||||
capture = ScreenshotCapture()
|
||||
async def test_is_available_when_no_engine(self):
|
||||
"""Test availability when no engine is available."""
|
||||
capture = ScreenshotCapture(ScreenshotConfig(enabled=True))
|
||||
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.returncode = 1
|
||||
mock_proc.wait = AsyncMock(return_value=1)
|
||||
mock_exec.return_value = mock_proc
|
||||
|
||||
result = await capture.is_available()
|
||||
assert result is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_available_handles_exception(self):
|
||||
"""Test availability handles exceptions gracefully."""
|
||||
capture = ScreenshotCapture()
|
||||
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_exec.side_effect = Exception("Process error")
|
||||
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
|
||||
|
|
@ -141,7 +134,7 @@ class TestScreenshotCaptureMethod:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capture_unavailable_returns_none(self):
|
||||
"""Test capture returns None when uri2png unavailable."""
|
||||
"""Test capture returns None when no engine available."""
|
||||
config = ScreenshotConfig(enabled=True)
|
||||
capture = ScreenshotCapture(config=config)
|
||||
|
||||
|
|
@ -149,72 +142,23 @@ class TestScreenshotCaptureMethod:
|
|||
result = await capture.capture("https://example.com")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capture_success(self):
|
||||
"""Test successful screenshot capture."""
|
||||
config = ScreenshotConfig(enabled=True)
|
||||
capture = ScreenshotCapture(config=config)
|
||||
|
||||
# Create a temp PNG file to simulate screenshot
|
||||
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
|
||||
# Write some fake PNG data
|
||||
fake_png = b'\x89PNG\r\n\x1a\n' + b'fake png data'
|
||||
tmp.write(fake_png)
|
||||
tmp_path = tmp.name
|
||||
|
||||
try:
|
||||
with patch.object(capture, 'is_available', new_callable=AsyncMock, return_value=True):
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.returncode = 0
|
||||
mock_proc.communicate = AsyncMock(return_value=(b'', b''))
|
||||
mock_exec.return_value = mock_proc
|
||||
|
||||
with patch('tempfile.NamedTemporaryFile') as mock_temp:
|
||||
mock_temp.return_value.__enter__ = Mock(return_value=Mock(name=tmp_path))
|
||||
mock_temp.return_value.__exit__ = Mock(return_value=False)
|
||||
|
||||
# Can't fully test without uri2png, but verify structure
|
||||
result = await capture.capture("https://example.com")
|
||||
# Result may be None if file checks fail, but no exception
|
||||
finally:
|
||||
if os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capture_timeout(self):
|
||||
"""Test capture handles timeout."""
|
||||
config = ScreenshotConfig(enabled=True)
|
||||
config = ScreenshotConfig(enabled=True, timeout=1000)
|
||||
capture = ScreenshotCapture(config=config)
|
||||
capture._initialized = True
|
||||
capture._engine = Mock()
|
||||
capture._engine_name = 'test'
|
||||
|
||||
with patch.object(capture, 'is_available', new_callable=AsyncMock, return_value=True):
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.communicate = AsyncMock(side_effect=asyncio.TimeoutError())
|
||||
mock_proc.kill = Mock()
|
||||
mock_proc.wait = AsyncMock()
|
||||
mock_exec.return_value = mock_proc
|
||||
# Make engine.capture raise TimeoutError
|
||||
async def slow_capture(*args):
|
||||
raise asyncio.TimeoutError()
|
||||
|
||||
with patch('tempfile.NamedTemporaryFile'):
|
||||
result = await capture.capture("https://slow-site.com")
|
||||
assert result is None
|
||||
capture._engine.capture = slow_capture
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_capture_process_failure(self):
|
||||
"""Test capture handles process failure."""
|
||||
config = ScreenshotConfig(enabled=True)
|
||||
capture = ScreenshotCapture(config=config)
|
||||
|
||||
with patch.object(capture, 'is_available', new_callable=AsyncMock, return_value=True):
|
||||
with patch('asyncio.create_subprocess_exec', new_callable=AsyncMock) as mock_exec:
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.returncode = 1
|
||||
mock_proc.communicate = AsyncMock(return_value=(b'', b'Error'))
|
||||
mock_exec.return_value = mock_proc
|
||||
|
||||
with patch('tempfile.NamedTemporaryFile'):
|
||||
result = await capture.capture("https://error-site.com")
|
||||
assert result is None
|
||||
result = await capture.capture("https://slow-site.com")
|
||||
assert result is None
|
||||
|
||||
|
||||
class TestScreenshotCaptureToFile:
|
||||
|
|
@ -280,7 +224,7 @@ class TestScreenshotResultStructure:
|
|||
|
||||
def test_expected_result_keys(self):
|
||||
"""Document expected keys in capture result."""
|
||||
expected_keys = ['data', 'md5_hash', 'mime_type', 'size', 'source_uri']
|
||||
expected_keys = ['data', 'md5_hash', 'mime_type', 'size', 'source_uri', 'engine']
|
||||
|
||||
# Just document the structure - actual capture requires uri2png
|
||||
result = {
|
||||
|
|
@ -289,11 +233,53 @@ class TestScreenshotResultStructure:
|
|||
'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'])
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue