test: print to pdf: convert to PNG and compare to golden

tests/screenshot/test_screenshot.py

@@ -13,41 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import base64
-import io
 from pathlib import Path
 
 import pytest
 from anys import ANY_STR
-from PIL import Image, ImageChops
-from test_helpers import (execute_command, get_tree, goto_url,
-                          read_JSON_message, send_JSON_command)
-
-
-def assert_images_equal(img1: Image, img2: Image):
-    """Assert that the given images are equal."""
-    equal_size = (img1.height == img2.height) and (img1.width == img2.width)
-
-    if img1.mode == img2.mode == "RGBA":
-        img1_alphas = [pixel[3] for pixel in img1.getdata()]
-        img2_alphas = [pixel[3] for pixel in img2.getdata()]
-        equal_alphas = img1_alphas == img2_alphas
-    else:
-        equal_alphas = True
-
-    equal_content = not ImageChops.difference(img1.convert("RGB"),
-                                              img2.convert("RGB")).getbbox()
-
-    assert equal_alphas
-    assert equal_size
-    assert equal_content
-
-
-def save_png(png_bytes_or_str: bytes | str, output_file: str):
-    """Save the given PNG (bytes or base64 string representation) to the given output file."""
-    png_bytes = png_bytes_or_str if isinstance(
-        png_bytes_or_str, bytes) else base64.b64decode(png_bytes_or_str,
-                                                       validate=True)
-    Image.open(io.BytesIO(png_bytes)).save(output_file, 'PNG')
+from test_helpers import (assert_images_equal, execute_command, get_tree,
+                          goto_url, read_JSON_message, send_JSON_command)
 
 
 @pytest.mark.asyncio
@@ -94,9 +65,7 @@ async def test_screenshot(websocket, context_id, png_filename,
         resp = await read_JSON_message(websocket)
         assert resp["result"] == {'data': ANY_STR}
 
-        img1 = Image.open(io.BytesIO(base64.b64decode(resp["result"]["data"])))
-        img2 = Image.open(io.BytesIO(base64.b64decode(png_base64)))
-        assert_images_equal(img1, img2)
+        assert_images_equal(resp["result"]["data"], png_base64)
 
 
 @pytest.mark.asyncio
@@ -146,6 +115,4 @@ async def test_screenshot_oopif(websocket, context_id, html, iframe,
               'rb') as image_file:
         png_base64 = base64.b64encode(image_file.read()).decode('utf-8')
 
-        img1 = Image.open(io.BytesIO(base64.b64decode(resp["result"]["data"])))
-        img2 = Image.open(io.BytesIO(base64.b64decode(png_base64)))
-        assert_images_equal(img1, img2)
+        assert_images_equal(resp["result"]["data"], png_base64)

tests/test_helpers.py

@@ -14,11 +14,14 @@
 # limitations under the License.
 from __future__ import annotations
 
+import base64
+import io
 import itertools
 import json
 from typing import Literal
 
 from anys import ANY_NUMBER, ANY_STR, AnyContains, AnyGT, AnyLT, AnyWithEntries
+from PIL import Image, ImageChops
 
 _command_counter = itertools.count(1)
 
@@ -188,3 +191,46 @@ def AnyExtending(expected: list | dict):
         return AnyWithEntries(dict_result)
 
     return expected
+
+
+def assert_images_equal(img1: Image.Image | str, img2: Image.Image | str):
+    """Assert that the given images are equal."""
+    if isinstance(img1, str):
+        img1 = Image.open(io.BytesIO(base64.b64decode(img1)))
+    if isinstance(img2, str):
+        img2 = Image.open(io.BytesIO(base64.b64decode(img2)))
+
+    equal_size = (img1.height == img2.height) and (img1.width == img2.width)
+
+    if img1.mode == img2.mode == "RGBA":
+        img1_alphas = [pixel[3] for pixel in img1.getdata()]
+        img2_alphas = [pixel[3] for pixel in img2.getdata()]
+        equal_alphas = img1_alphas == img2_alphas
+    else:
+        equal_alphas = True
+
+    equal_content = not ImageChops.difference(img1.convert("RGB"),
+                                              img2.convert("RGB")).getbbox()
+
+    assert equal_alphas
+    assert equal_size
+    assert equal_content
+
+
+def save_png(png_bytes_or_str: bytes | str, output_file: str):
+    """Save the given PNG (bytes or base64 string representation) to the given output file."""
+    png_bytes = png_bytes_or_str if isinstance(
+        png_bytes_or_str, bytes) else base64.b64decode(png_bytes_or_str,
+                                                       validate=True)
+    Image.open(io.BytesIO(png_bytes)).save(output_file, 'PNG')
+
+
+def save_pdf(pdf_bytes_or_str: bytes | str, output_file: str):
+    pdf_bytes = pdf_bytes_or_str if isinstance(
+        pdf_bytes_or_str, bytes) else base64.b64decode(pdf_bytes_or_str,
+                                                       validate=True)
+    if pdf_bytes[0:4] != b'%PDF':
+        raise ValueError('Missing the PDF file signature')
+
+    with open(output_file, 'wb') as f:
+        f.write(pdf_bytes)

tests/test_print.py

@@ -13,29 +13,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import base64
-
 import pytest
 from anys import ANY_STR
-from test_helpers import execute_command, goto_url
-
-
-def save_pdf(pdf_bytes_or_str: bytes | str, output_file: str):
-    pdf_bytes = pdf_bytes_or_str if isinstance(
-        pdf_bytes_or_str, bytes) else base64.b64decode(pdf_bytes_or_str,
-                                                       validate=True)
-    if pdf_bytes[0:4] != b'%PDF':
-        raise ValueError('Missing the PDF file signature')
-
-    with open(output_file, 'wb') as f:
-        f.write(pdf_bytes)
+from test_helpers import assert_images_equal, execute_command, goto_url
 
 
 @pytest.mark.asyncio
-async def test_print(websocket, context_id, html):
+async def test_print(websocket, context_id, html, get_cdp_session_id):
     await goto_url(websocket, context_id, html())
 
-    result = await execute_command(
+    print_result = await execute_command(
         websocket, {
             "method": "browsingContext.print",
             "params": {
@@ -49,4 +36,45 @@ async def test_print(websocket, context_id, html):
         })
 
     # 'data' is not deterministic, ~a dozen characters differ between runs.
-    assert result['data'] == ANY_STR
+    assert print_result["data"] == ANY_STR
+
+    try:
+        await goto_url(websocket, context_id,
+                       f'data:application/pdf,base64;{print_result["data"]}')
+    except Exception as e:
+        assert e.args[0] == {
+            'error': 'unknown error',
+            'message': 'net::ERR_ABORTED'
+        }
+        pytest.xfail("PDF viewer not available in headless.")
+
+    session_id = await get_cdp_session_id(context_id)
+
+    # Set a fixed viewport to make the test deterministic.
+    await execute_command(
+        websocket, {
+            "method": "cdp.sendCommand",
+            "params": {
+                "method": "Emulation.setDeviceMetricsOverride",
+                "params": {
+                    "width": 200,
+                    "height": 200,
+                    "deviceScaleFactor": 1.0,
+                    "mobile": False,
+                },
+                "session": session_id
+            }
+        })
+
+    screenshot_result = await execute_command(
+        websocket, {
+            "method": "browsingContext.captureScreenshot",
+            "params": {
+                "context": context_id
+            }
+        })
+
+    assert_images_equal(
+        screenshot_result["data"],
+        "iVBORw0KGgoAAAANSUhEUgAAAMgAAADICAYAAACtWK6eAAAAAXNSR0IArs4c6QAAAhlJREFUeJzt07ENgDAAwLDSEzsg8f8h9IEqKwz2BVlyrft5B3A0vw6APzMIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQDAIBINAMAgEg0AwCASDQNhGJAOQPl8yRgAAAABJRU5ErkJggg=="
+    )