From 36edc3b8dc34f1c9f8ea308ce44af306c80856f3 Mon Sep 17 00:00:00 2001 From: Russell Ballestrini Date: Mon, 29 Dec 2025 11:35:14 -0500 Subject: [PATCH] modified: archive.py modified: tests/unit/test_screenshot.py --- archive.py | 30 +++++- tests/unit/test_screenshot.py | 198 ++++++++++++++++------------------ 2 files changed, 117 insertions(+), 111 deletions(-) diff --git a/archive.py b/archive.py index 64fb0a3..e734580 100644 --- a/archive.py +++ b/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, '', '', '...', 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( diff --git a/tests/unit/test_screenshot.py b/tests/unit/test_screenshot.py index eefd0d3..ebba280 100644 --- a/tests/unit/test_screenshot.py +++ b/tests/unit/test_screenshot.py @@ -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'])