test: add tests to new take_screenshot behavior

Author: thalissonvs

tests/test_browser/test_browser_tab.py

@@ -464,7 +464,7 @@ async def test_print_to_pdf_as_base64(self, tab):
             'result': {'data': pdf_data}
         }
 
-        result = await tab.print_to_pdf('', as_base64=True)
+        result = await tab.print_to_pdf(as_base64=True)
 
         assert result == pdf_data
         assert_mock_called_at_least_once(tab._connection_handler)
@@ -1779,17 +1779,15 @@ async def test_take_screenshot_invalid_extension(self, tab):
 
     @pytest.mark.asyncio
     async def test_print_to_pdf_with_invalid_path(self, tab):
-        """Test print_to_pdf with invalid path handling."""
+        """Test print_to_pdf with missing path when not using base64."""
         # Mock the response
         tab._connection_handler.execute_command.return_value = {
             'result': {'data': 'JVBERi0xLjQKJdPr6eEKMSAwIG9iago8PAovVHlwZSAvQ2F0YWxvZwo+PgplbmRvYmoKdHJhaWxlcgo8PAovUm9vdCAxIDAgUgo+PgpzdGFydHhyZWYKMTgKJSVFT0Y='}
         }
 
-        # Should not raise exception - print_to_pdf doesn't validate extensions
-        result = await tab.print_to_pdf('document.txt', as_base64=True)
-        
-        assert result is not None
-        assert_mock_called_at_least_once(tab._connection_handler)
+        # Should raise ValueError when path is not provided and as_base64=False
+        with pytest.raises(ValueError, match="path is required when as_base64=False"):
+            await tab.print_to_pdf(as_base64=False)
 
     @pytest.mark.asyncio
     async def test_execute_script_with_none_element(self, tab):

tests/test_web_element.py

@@ -631,7 +631,7 @@ async def test_take_screenshot_success(self, web_element, tmp_path):
             {'result': {'data': screenshot_data}},  # capture_screenshot
         ]
 
-        screenshot_path = tmp_path / 'element.jpg'
+        screenshot_path = tmp_path / 'element.jpeg'
 
         # Mock aiofiles.open properly for async context manager
         mock_file = AsyncMock()
@@ -654,7 +654,7 @@ async def test_take_screenshot_default_quality(self, web_element, tmp_path):
             {'result': {'data': screenshot_data}},
         ]
 
-        screenshot_path = tmp_path / 'element_default.jpg'
+        screenshot_path = tmp_path / 'element_default.jpeg'
 
         # Mock aiofiles.open properly for async context manager
         mock_file = AsyncMock()
@@ -667,6 +667,57 @@ async def test_take_screenshot_default_quality(self, web_element, tmp_path):
         # Should call get_bounds_using_js and capture_screenshot
         assert web_element._connection_handler.execute_command.call_count == 2
 
+    @pytest.mark.asyncio
+    async def test_take_screenshot_as_base64(self, web_element):
+        """Test screenshot returned as base64 string."""
+        bounds = {'x': 10, 'y': 20, 'width': 100, 'height': 50}
+        screenshot_data = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/wcAAgAB/edzE+oAAAAASUVORK5CYII='
+
+        web_element._connection_handler.execute_command.side_effect = [
+            {'result': {'result': {'value': json.dumps(bounds)}}},  # get_bounds_using_js
+            {'result': {'data': screenshot_data}},  # capture_screenshot
+        ]
+
+        # Take screenshot as base64
+        result = await web_element.take_screenshot(as_base64=True)
+
+        # Should return the base64 data
+        assert result == screenshot_data
+        # Should call get_bounds_using_js and capture_screenshot
+        assert web_element._connection_handler.execute_command.call_count == 2
+
+    @pytest.mark.asyncio
+    async def test_take_screenshot_missing_path_without_base64(self, web_element):
+        """Test screenshot raises error when no path and as_base64=False."""
+        from pydoll.exceptions import MissingScreenshotPath
+
+        with pytest.raises(MissingScreenshotPath):
+            await web_element.take_screenshot(as_base64=False)
+
+    @pytest.mark.asyncio
+    async def test_take_screenshot_jpg_alias(self, web_element, tmp_path):
+        """Test that .jpg extension works as alias for .jpeg."""
+        bounds = {'x': 10, 'y': 20, 'width': 100, 'height': 50}
+        screenshot_data = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/wcAAgAB/edzE+oAAAAASUVORK5CYII='
+
+        web_element._connection_handler.execute_command.side_effect = [
+            {'result': {'result': {'value': json.dumps(bounds)}}},  # get_bounds_using_js
+            {'result': {'data': screenshot_data}},  # capture_screenshot
+        ]
+
+        screenshot_path = tmp_path / 'element.jpg'
+
+        # Mock aiofiles.open properly for async context manager
+        mock_file = AsyncMock()
+        mock_file.write = AsyncMock()
+
+        with patch('aiofiles.open') as mock_aiofiles_open:
+            mock_aiofiles_open.return_value.__aenter__.return_value = mock_file
+            await web_element.take_screenshot(str(screenshot_path), quality=90)
+
+        # Should work without raising InvalidFileExtension
+        assert web_element._connection_handler.execute_command.call_count == 2
+
 
 class TestWebElementVisibility:
     """Test element visibility and interaction checks."""