feat: full-page screenshots (firefox raw command, chromium PIL stitch)

2026-07-07 23:14:53 -06:00
parent 2adb7b2236
commit ab50a984a8
2 changed files with 140 additions and 1 deletions
@@ -0,0 +1,68 @@
+import base64
+from io import BytesIO
+from unittest.mock import MagicMock
+
+from PIL import Image
+
+from wabot.screenshot import FULL_PAGE_SCREENSHOT, save_full_page
+
+
+def png_bytes(width, height, color=(200, 30, 30)):
+    buf = BytesIO()
+    Image.new("RGB", (width, height), color).save(buf, format="PNG")
+    return buf.getvalue()
+
+
+class TestFirefoxPath:
+    def test_registers_raw_command_and_decodes_base64(self, tmp_path):
+        driver = MagicMock(name="driver")
+        driver.capabilities = {"browserName": "firefox"}
+        driver.command_executor._commands = {}
+        driver.execute.return_value = {"value": base64.b64encode(png_bytes(50, 80)).decode()}
+        out = tmp_path / "shot.png"
+        assert save_full_page(driver, str(out)) is True
+        # helper methods only exist on webdriver.Firefox, so the raw
+        # geckodriver command must have been registered for Remote support
+        assert driver.command_executor._commands[FULL_PAGE_SCREENSHOT] == (
+            "GET",
+            "/session/$sessionId/moz/screenshot/full",
+        )
+        driver.execute.assert_called_once_with(FULL_PAGE_SCREENSHOT)
+        assert Image.open(out).size == (50, 80)
+
+
+class TestChromiumPath:
+    def test_single_viewport_page_is_saved_directly(self, tmp_path):
+        driver = MagicMock(name="driver")
+        driver.capabilities = {"browserName": "chrome"}
+        # page fits in one viewport: total == viewport
+        driver.execute_script.side_effect = lambda script: {
+            "return document.body.parentNode.scrollWidth": 100,
+            "return document.body.parentNode.scrollHeight": 60,
+            "return document.documentElement.clientWidth": 100,
+            "return window.innerHeight": 60,
+        }.get(script)
+        driver.get_screenshot_as_png.return_value = png_bytes(100, 60)
+        out = tmp_path / "shot.png"
+        assert save_full_page(driver, str(out)) is True
+        assert Image.open(out).size == (100, 60)
+
+    def test_tall_page_is_stitched_from_tiles(self, tmp_path):
+        driver = MagicMock(name="driver")
+        driver.capabilities = {"browserName": "chrome"}
+        driver.execute_script.side_effect = lambda script: {
+            "return document.body.parentNode.scrollWidth": 100,
+            "return document.body.parentNode.scrollHeight": 150,  # 3 tiles of 60
+            "return document.documentElement.clientWidth": 100,
+            "return window.innerHeight": 60,
+        }.get(script)
+        driver.get_screenshot_as_png.return_value = png_bytes(100, 60)
+        out = tmp_path / "shot.png"
+        assert save_full_page(driver, str(out)) is True
+        assert Image.open(out).size == (100, 150)
+        # scrolled at least twice beyond the initial position
+        scroll_calls = [
+            c for c in driver.execute_script.call_args_list
+            if c.args and str(c.args[0]).startswith("window.scrollTo")
+        ]
+        assert len(scroll_calls) >= 2