mirror of
https://git.zavage.net/Zavage-Software/wabot.git
synced 2026-07-21 13:06:08 -06:00
111 lines
5.1 KiB
Python
111 lines
5.1 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
|
|
|
|
def test_tall_page_stitches_tiles_bottom_aligned(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) # window.scrollTo(...) -> None
|
|
driver.get_screenshot_as_png.side_effect = [
|
|
png_bytes(100, 60, (255, 0, 0)), # tile @ y=0
|
|
png_bytes(100, 60, (0, 255, 0)), # tile @ y=60
|
|
png_bytes(100, 60, (0, 0, 255)), # tile @ y=120, clamped to paste_y=90
|
|
]
|
|
out = tmp_path / "shot.png"
|
|
assert save_full_page(driver, str(out)) is True
|
|
img = Image.open(out)
|
|
assert img.size == (100, 150)
|
|
assert img.getpixel((50, 30)) == (255, 0, 0) # first tile
|
|
assert img.getpixel((50, 75)) == (0, 255, 0) # second tile (rows 60-89)
|
|
assert img.getpixel((50, 140)) == (0, 0, 255) # third tile, bottom-aligned
|
|
|
|
def test_wide_page_clamps_last_column(self, tmp_path):
|
|
driver = MagicMock(name="driver")
|
|
driver.capabilities = {"browserName": "chrome"}
|
|
driver.execute_script.side_effect = lambda script: {
|
|
"return document.body.parentNode.scrollWidth": 150, # 2 columns of 100
|
|
"return document.body.parentNode.scrollHeight": 60,
|
|
"return document.documentElement.clientWidth": 100,
|
|
"return window.innerHeight": 60,
|
|
}.get(script)
|
|
driver.get_screenshot_as_png.side_effect = [
|
|
png_bytes(100, 60, (255, 0, 0)), # column @ x=0
|
|
png_bytes(100, 60, (0, 0, 255)), # column @ x=100, clamped to paste_x=50
|
|
]
|
|
out = tmp_path / "shot.png"
|
|
assert save_full_page(driver, str(out)) is True
|
|
img = Image.open(out)
|
|
assert img.size == (150, 60)
|
|
assert img.getpixel((25, 30)) == (255, 0, 0) # left column (cols 0-49)
|
|
assert img.getpixel((140, 30)) == (0, 0, 255) # right column, clamped (cols 50-149)
|