tests/unit_tests/test_cli_screenshot.py

0001 # test_cli_screenshot.py
0002
0003 import pytest
0004 import os
0005 import sys
0006 from unittest.mock import patch, MagicMock
0007 from pyrobird.cli.screenshot import get_screenshot_path, capture_screenshot
0008
0009 # Check if playwright is available
0010 try:
0011     import playwright.sync_api
0012     PLAYWRIGHT_AVAILABLE = True
0013 except ImportError:
0014     PLAYWRIGHT_AVAILABLE = False
0015
0016
0017 @pytest.fixture
0018 def temp_screenshots_dir(tmp_path):
0019     """Create a temporary screenshots directory."""
0020     screenshots_dir = tmp_path / "screenshots"
0021     screenshots_dir.mkdir()
0022     return tmp_path
0023
0024
0025 def test_get_screenshot_path_creates_directory(tmp_path):
0026     """Test that get_screenshot_path creates the screenshots directory."""
0027     # Change to temp directory
0028     original_cwd = os.getcwd()
0029     os.chdir(tmp_path)
0030
0031     try:
0032         output_path = get_screenshot_path("test.png")
0033         assert os.path.exists("screenshots"), "Screenshots directory was not created"
0034         assert output_path == os.path.join("screenshots", "test.png")
0035     finally:
0036         os.chdir(original_cwd)
0037
0038
0039 def test_get_screenshot_path_no_extension(tmp_path):
0040     """Test that get_screenshot_path adds .png extension when none is provided."""
0041     original_cwd = os.getcwd()
0042     os.chdir(tmp_path)
0043
0044     try:
0045         output_path = get_screenshot_path("test")
0046         assert output_path == os.path.join("screenshots", "test.png")
0047     finally:
0048         os.chdir(original_cwd)
0049
0050
0051 def test_get_screenshot_path_auto_numbering(tmp_path):
0052     """Test that get_screenshot_path handles auto-numbering correctly."""
0053     original_cwd = os.getcwd()
0054     os.chdir(tmp_path)
0055
0056     try:
0057         screenshots_dir = tmp_path / "screenshots"
0058         screenshots_dir.mkdir()
0059
0060         # Create existing screenshot files
0061         (screenshots_dir / "test.png").touch()
0062         (screenshots_dir / "test_001.png").touch()
0063         (screenshots_dir / "test_002.png").touch()
0064
0065         # Next screenshot should be test_003.png
0066         output_path = get_screenshot_path("test.png")
0067         assert output_path == os.path.join("screenshots", "test_003.png")
0068     finally:
0069         os.chdir(original_cwd)
0070
0071
0072 def test_get_screenshot_path_no_existing_files(tmp_path):
0073     """Test get_screenshot_path when no files exist yet."""
0074     original_cwd = os.getcwd()
0075     os.chdir(tmp_path)
0076
0077     try:
0078         output_path = get_screenshot_path("screenshot.png")
0079         assert output_path == os.path.join("screenshots", "screenshot.png")
0080     finally:
0081         os.chdir(original_cwd)
0082
0083
0084 def test_get_screenshot_path_with_gaps_in_numbering(tmp_path):
0085     """Test that get_screenshot_path finds the next number even with gaps."""
0086     original_cwd = os.getcwd()
0087     os.chdir(tmp_path)
0088
0089     try:
0090         screenshots_dir = tmp_path / "screenshots"
0091         screenshots_dir.mkdir()
0092
0093         # Create files with gaps in numbering
0094         (screenshots_dir / "test.png").touch()
0095         (screenshots_dir / "test_001.png").touch()
0096         (screenshots_dir / "test_005.png").touch()  # Gap in numbering
0097
0098         # Next screenshot should be test_006.png (max + 1)
0099         output_path = get_screenshot_path("test.png")
0100         assert output_path == os.path.join("screenshots", "test_006.png")
0101     finally:
0102         os.chdir(original_cwd)
0103
0104
0105 def test_capture_screenshot_missing_playwright():
0106     """Test that capture_screenshot handles missing playwright gracefully."""
0107     with patch('builtins.__import__', side_effect=ImportError("No module named 'playwright'")):
0108         with pytest.raises(SystemExit) as exc_info:
0109             capture_screenshot("http://example.com", "/tmp/test.png")
0110         assert exc_info.value.code == 1
0111
0112
0113 @pytest.mark.skipif(not PLAYWRIGHT_AVAILABLE, reason="playwright not installed")
0114 @patch('playwright.sync_api.sync_playwright')
0115 def test_capture_screenshot_basic_workflow(mock_sync_playwright, tmp_path):
0116     """Test the basic workflow of capture_screenshot."""
0117     # Setup mocks
0118     mock_playwright = MagicMock()
0119     mock_browser = MagicMock()
0120     mock_page = MagicMock()
0121
0122     mock_sync_playwright.return_value.__enter__.return_value = mock_playwright
0123     mock_playwright.chromium.launch.return_value = mock_browser
0124     mock_browser.new_page.return_value = mock_page
0125
0126     # Test URL and output path
0127     test_url = "http://localhost:5454"
0128     output_file = str(tmp_path / "test_screenshot.png")
0129
0130     # Call the function
0131     capture_screenshot(test_url, output_file)
0132
0133     # Verify the workflow
0134     mock_playwright.chromium.launch.assert_called_once_with(headless=True)
0135     mock_browser.new_page.assert_called_once()
0136     mock_page.set_viewport_size.assert_called_once_with({"width": 1920, "height": 1080})
0137     mock_page.goto.assert_called_once_with(test_url)
0138     mock_page.screenshot.assert_called_once_with(path=output_file, full_page=True)
0139     mock_browser.close.assert_called_once()
0140
0141
0142 @pytest.mark.skipif(not PLAYWRIGHT_AVAILABLE, reason="playwright not installed")
0143 @patch('playwright.sync_api.sync_playwright')
0144 def test_capture_screenshot_wait_for_load_state(mock_sync_playwright, tmp_path):
0145     """Test that capture_screenshot waits for the page to load."""
0146     # Setup mocks
0147     mock_playwright = MagicMock()
0148     mock_browser = MagicMock()
0149     mock_page = MagicMock()
0150
0151     mock_sync_playwright.return_value.__enter__.return_value = mock_playwright
0152     mock_playwright.chromium.launch.return_value = mock_browser
0153     mock_browser.new_page.return_value = mock_page
0154
0155     test_url = "http://localhost:5454"
0156     output_file = str(tmp_path / "test_screenshot.png")
0157
0158     # Call the function
0159     capture_screenshot(test_url, output_file)
0160
0161     # Verify wait_for_load_state was called
0162     mock_page.wait_for_load_state.assert_called_once_with("domcontentloaded", timeout=10_000)
0163
0164
0165 @pytest.mark.skipif(not PLAYWRIGHT_AVAILABLE, reason="playwright not installed")
0166 @patch('playwright.sync_api.sync_playwright')
0167 @patch('pyrobird.cli.screenshot.time.sleep')
0168 def test_capture_screenshot_fallback_to_selector(mock_sleep, mock_sync_playwright, tmp_path):
0169     """Test that capture_screenshot falls back to wait_for_selector if wait_for_load_state fails."""
0170     # Setup mocks
0171     mock_playwright = MagicMock()
0172     mock_browser = MagicMock()
0173     mock_page = MagicMock()
0174
0175     mock_sync_playwright.return_value.__enter__.return_value = mock_playwright
0176     mock_playwright.chromium.launch.return_value = mock_browser
0177     mock_browser.new_page.return_value = mock_page
0178
0179     # Make wait_for_load_state raise an exception
0180     mock_page.wait_for_load_state.side_effect = Exception("Timeout")
0181
0182     test_url = "http://localhost:5454"
0183     output_file = str(tmp_path / "test_screenshot.png")
0184
0185     # Call the function
0186     capture_screenshot(test_url, output_file)
0187
0188     # Verify fallback to wait_for_selector was attempted
0189     mock_page.wait_for_selector.assert_called_once_with('body', timeout=10_000)
0190
0191
0192 @pytest.mark.skipif(not PLAYWRIGHT_AVAILABLE, reason="playwright not installed")
0193 @patch('playwright.sync_api.sync_playwright')
0194 @patch('pyrobird.cli.screenshot.time.sleep')
0195 def test_capture_screenshot_double_fallback(mock_sleep, mock_sync_playwright, tmp_path):
0196     """Test that capture_screenshot falls back to sleep if both wait methods fail."""
0197     # Setup mocks
0198     mock_playwright = MagicMock()
0199     mock_browser = MagicMock()
0200     mock_page = MagicMock()
0201
0202     mock_sync_playwright.return_value.__enter__.return_value = mock_playwright
0203     mock_playwright.chromium.launch.return_value = mock_browser
0204     mock_browser.new_page.return_value = mock_page
0205
0206     # Make both wait methods raise exceptions
0207     mock_page.wait_for_load_state.side_effect = Exception("Timeout")
0208     mock_page.wait_for_selector.side_effect = Exception("Timeout")
0209
0210     test_url = "http://localhost:5454"
0211     output_file = str(tmp_path / "test_screenshot.png")
0212
0213     # Call the function
0214     capture_screenshot(test_url, output_file)
0215
0216     # Verify sleep was called as the final fallback
0217     assert mock_sleep.call_count >= 1  # At least one sleep(3) call in the fallback