feat: add on-demand device screenshots

2026-08-31 10:43:37 +08:00
parent 60ee157e97
commit 8c99dc015a
10 changed files with 231 additions and 2 deletions
@@ -242,6 +242,78 @@ def test_add_device_appears_in_devices_page_and_manager(tmp_path) -> None:
    assert [device.id for device in context["manager"].list_devices()] == ["device-a"]


+def test_devices_page_captures_screenshot_only_when_button_endpoint_is_called(
+    tmp_path,
+) -> None:
+    class ScreenshotDriver:
+        def __init__(self) -> None:
+            self.capture_count = 0
+
+        def connect(self) -> None:
+            pass
+
+        def disconnect(self) -> None:
+            pass
+
+        def screenshot(self) -> bytes:
+            self.capture_count += 1
+            return b"fake-png"
+
+    driver = ScreenshotDriver()
+    client, context = _build_client(tmp_path)
+    context["config_store"].add(
+        device_id="device-a",
+        name="Lab iPhone",
+        driver_type="wda",
+        connection_info={},
+    )
+    context["manager"].register_device("device-a", lambda: driver)
+    context["manager"].connect("device-a")
+    csrf_token = _login(client)
+
+    page = client.get("/devices")
+    assert page.status_code == 200
+    assert 'class="screenshot-button"' in page.text
+    assert 'data-device-id="device-a"' in page.text
+    assert driver.capture_count == 0
+
+    response = client.post(
+        "/api/devices/device-a/screenshot",
+        headers={"X-CSRF-Token": csrf_token},
+    )
+
+    assert response.status_code == 200
+    assert response.content == b"fake-png"
+    assert response.headers["content-type"] == "image/png"
+    assert response.headers["cache-control"] == "no-store"
+    assert driver.capture_count == 1
+
+
+def test_device_screenshot_requires_csrf_and_connected_device(tmp_path) -> None:
+    class ScreenshotDriver:
+        def connect(self) -> None:
+            pass
+
+        def disconnect(self) -> None:
+            pass
+
+        def screenshot(self) -> bytes:
+            return b"fake-png"
+
+    client, context = _build_client(tmp_path)
+    context["manager"].register_device("device-a", ScreenshotDriver)
+    csrf_token = _login(client)
+
+    missing_csrf = client.post("/api/devices/device-a/screenshot")
+    assert missing_csrf.status_code == 403
+
+    offline = client.post(
+        "/api/devices/device-a/screenshot",
+        headers={"X-CSRF-Token": csrf_token},
+    )
+    assert offline.status_code == 503
+
+
 def test_remove_device_unregisters_from_manager(tmp_path) -> None:
    client, context = _build_client(tmp_path)
    csrf_token = _login(client)