mirror of
https://git.zavage.net/Zavage-Software/wabot.git
synced 2026-07-21 13:06:08 -06:00
69 lines
2.9 KiB
Python
69 lines
2.9 KiB
Python
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
|