Files
ha-nvidia-shield/tests/test_screen_capture_image.py
T
James Mackay f29f62ddae feat(image): opt-in screen capture as an image entity, binary-safe over ADB
The user asked for "an image entity of the device screen capture? Can that be
added". It is added as `image.<entry>_screen_capture`, and the three boundaries
are the feature:

- it NEVER drives state — the state is the capture timestamp, and the capture
  path reads the ADB poller and nothing else (pinned by a manager that raises on
  any other attribute, and by a real-HA run whose recorded state is the
  timestamp with no bytes anywhere in the attributes);
- it is OPT-IN and OFF by default (`screen_capture_enabled`, Main setup row) —
  with the option absent `async_setup_entry` adds no entity at all;
- it NEVER polls — HA's own `should_poll = False` for image entities is kept, so
  only an image fetch or `homeassistant.update_entity` captures; concurrent
  fetches share one shell and a capture is reused for 5 s.

BINARY-SAFE CAPTURE: `screencap -p | base64` over the existing shell, decoded
locally with `base64.b64decode` and gated on the PNG signature AND the IHDR
chunk. `AdbDevice.shell()` returns `str`, so raw PNG bytes cannot travel through
it; the alternative (a device temp file + the library's file-read path) needs the
raw AdbDevice and writes to the Shield's shared storage, where a crash between
write and delete leaves a screenshot behind. The base64 route leaves no
device-side artifact — nothing to clean up — at the cost of ~33% more bytes on
the wire.

HONEST FAILURES: ADB off / ADB down / a capture that returned no PNG / a
transport that raised all make the entity UNAVAILABLE with `capture_status`
(`adb_disabled` / `adb_disconnected` / `capture_failed`) — never a blank or
zero-byte image. `image_last_updated` moves only on success.

Tests: 25 new cases (a real 2x1 PNG fixture on the wire — bytes, never a str;
every non-PNG answer rejected; disconnected/failed/recovery paths; capture reuse,
coalescing and the expiry window; opt-in default OFF and on) plus 4 config-flow
cases (schema default False, off stays off, on is stored, and it can be turned
back off). The stub harness gains the image platform and `util.dt` faces, the
real-HA entity matrix gains this class, and the drift sweep's coverage test
demanded exactly that. README documents the cost, the opt-in and the boundary.
2026-10-12 10:28:32 +01:00

384 lines
15 KiB
Python

"""t66: the opt-in screen-capture IMAGE entity.
WHAT IS PROVEN HERE, AND WHAT IS NOT
------------------------------------
Proven: the bytes that reach Home Assistant are REAL PNG bytes — ``PNG_B64``
below is a real 2x1 PNG, base64'd on the wire exactly the way
``screencap -p | base64`` sends it — the entity hands HA ``bytes`` and never a
``str``, a disconnected link or a bad answer makes the entity UNAVAILABLE
instead of serving a blank/zero-byte image, one capture is coalesced and reused
so a dashboard cannot hammer the device, and the option defaults OFF so no
entity exists until the user opts in.
NOT proven: anything about a real Shield. No device was available in this
environment, so ``screencap``'s flags, the device's base64 and — the important
one — how the persistent ADB socket copes with a few hundred kilobytes of
base64 in a single ``shell()`` call are INFERRED, not measured. ``FakeAdbPoller``
answers the way a working Shield should.
"""
import asyncio
import base64
import json
import sys
from datetime import datetime
from types import SimpleNamespace
import entity_stublib
import pytest
from conftest import _SHIELD_DIR
from entity_stublib import install, load_entity_modules
#: A REAL 2x1 PNG (one red pixel, one green pixel) — 72 bytes, base64. It is the
#: wire form the device sends, so the fixture exercises the real decode path
#: rather than an invented one.
PNG_B64 = (
"iVBORw0KGgoAAAANSUhEUgAAAAIAAAABCAIAAAB7QOjdAAAAD0lEQVR42mP4z8DA8J8BAAf/"
"Af+4BDXgAAAAAElFTkSuQmCC"
)
PNG_BYTES = base64.b64decode(PNG_B64)
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
def wrapped(text: str, width: int = 76) -> str:
"""``base64``'s own line wrapping — the device wraps, and that is not corruption."""
return "\n".join(text[i : i + width] for i in range(0, len(text), width)) + "\n"
@pytest.fixture(autouse=True)
def scoped_entity_stubs():
"""Container for the HA stubs + loaded entity modules (t7 discipline)."""
snapshot = entity_stublib.snapshot_modules()
install()
yield load_entity_modules()
entity_stublib.restore_modules(snapshot)
@pytest.fixture(name="modules")
def load_modules_fixture():
return load_entity_modules()
class FakeAdbPoller:
"""The ADB seam the entity is allowed to use: ``connected`` + ``shell``."""
def __init__(self, answer=None, *, connected=True, delay=0.0):
self.connected = connected
self.answer = answer
self.delay = delay
self.calls: list[str] = []
async def shell(self, cmd: str):
self.calls.append(cmd)
if self.delay:
await asyncio.sleep(self.delay)
if isinstance(self.answer, Exception):
raise self.answer
if callable(self.answer):
return self.answer(cmd)
return self.answer
class StrictManager:
"""Only ``adb_poller`` may be touched — nothing else exists on purpose.
"Never a state source" is a claim about what the capture path reads and
writes: it reads the ADB poller for a screen grab and writes nothing at all.
A manager that raises on every other attribute turns a future reach-in (a
store write, a state push) into a failing test instead of a quiet coupling
between screen capture and the state pipeline.
"""
def __init__(self, poller):
self.adb_poller = poller
def __getattr__(self, name):
raise AssertionError(f"the screen-capture path touched manager.{name}")
def make_entry(options=None, *, entry_id="e1", title="Shield"):
return SimpleNamespace(
entry_id=entry_id, title=title, options=dict(options or {}), runtime_data=None
)
def make_entity(modules, poller, options=None):
"""The entity over a fake manager on the entry, plus the entry it belongs to."""
entry = make_entry(options or {"screen_capture_enabled": True})
if poller is not None:
entry.runtime_data = StrictManager(poller)
entity = modules["image"].NvidiaShieldScreenImage(SimpleNamespace(), entry)
return entity, entry
def fake_clock(modules, monkeypatch, start=1000.0):
"""Freeze ``time.monotonic`` for the capture-window tests. Returns the cell."""
now = [start]
monkeypatch.setattr(
modules["image"], "time", SimpleNamespace(monotonic=lambda: now[0])
)
return now
# ── the decode path ──────────────────────────────────────────────────────────
class TestScreenCaptureDecode:
"""``decode_screen_capture`` — the one place an image could be corrupted."""
def test_a_real_png_survives_the_wire_format(self, modules):
decode = modules["image"].decode_screen_capture
data = decode(wrapped(PNG_B64))
assert data == PNG_BYTES # byte-for-byte, wrapping and all
assert data.startswith(PNG_SIGNATURE)
assert isinstance(data, bytes) # not a str: this is base64-decoded
@pytest.mark.parametrize(
"answer",
[
None, # the transport never answered (AdbClientBase.shell's failure)
"", # an empty answer is not an image
"not base64 at all !!",
base64.b64encode(b"hello world, definitely not a png").decode(),
# a signature and nothing else: not a PNG (the IHDR chunk is missing)
base64.b64encode(PNG_SIGNATURE + b"\x00" * 8).decode(),
],
)
def test_everything_that_is_not_a_png_is_rejected(self, modules, answer):
assert modules["image"].decode_screen_capture(answer) is None
def test_a_device_error_line_is_not_an_image(self, modules):
polluted = "error: no display\n" + wrapped(PNG_B64)
assert modules["image"].decode_screen_capture(polluted) is None
# ── the entity ───────────────────────────────────────────────────────────────
class TestScreenCaptureEntity:
async def test_the_entity_serves_real_png_bytes_and_never_a_str(self, modules):
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, _entry = make_entity(modules, poller)
data = await entity.async_image()
assert data == PNG_BYTES
assert data.startswith(PNG_SIGNATURE)
assert not isinstance(data, str)
assert entity.available is True
assert entity.capture_status == "ok"
assert entity.extra_state_attributes["capture_status"] == "ok"
async def test_the_command_is_the_base64_route_and_needs_no_device_file(self, modules):
"""Route 2 of the module docstring, and the reason it was chosen."""
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, _entry = make_entity(modules, poller)
await entity.async_image()
assert poller.calls == ["screencap -p | base64"]
# No device-side artifact exists, so there is nothing to clean up...
assert "/sdcard" not in poller.calls[0]
assert "rm " not in poller.calls[0]
async def test_a_disconnected_link_is_unavailable_and_is_never_asked(self, modules):
poller = FakeAdbPoller(wrapped(PNG_B64), connected=False)
entity, _entry = make_entity(modules, poller)
assert entity.available is False
assert entity.capture_status == "adb_disconnected"
assert await entity.async_image() is None
assert poller.calls == [] # a dead link is not poked for a screenshot
async def test_adb_disabled_has_its_own_reason(self, modules):
entry = make_entry({"screen_capture_enabled": True})
manager = SimpleNamespace(adb_poller=None)
entry.runtime_data = manager
entity = modules["image"].NvidiaShieldScreenImage(SimpleNamespace(), entry)
assert entity.available is False
assert entity.capture_status == "adb_disabled"
assert await entity.async_image() is None
async def test_an_empty_answer_is_unavailable_not_a_blank_image(self, modules):
poller = FakeAdbPoller(None)
entity, _entry = make_entity(modules, poller)
assert await entity.async_image() is None
assert entity.available is False
assert entity.capture_status == "capture_failed"
async def test_a_transport_that_raises_is_a_failure_not_a_crash(self, modules):
poller = FakeAdbPoller(RuntimeError("socket closed under us"))
entity, _entry = make_entity(modules, poller)
assert await entity.async_image() is None
assert entity.available is False
assert entity.capture_status == "capture_failed"
async def test_a_failed_capture_recovers_on_the_next_request(self, modules):
poller = FakeAdbPoller(None)
entity, _entry = make_entity(modules, poller)
assert await entity.async_image() is None
assert entity.available is False
poller.answer = wrapped(PNG_B64) # the device came back
assert await entity.async_image() == PNG_BYTES
assert entity.available is True
assert entity.capture_status == "ok"
async def test_one_capture_is_reused_inside_the_window(self, modules, monkeypatch):
"""A dashboard repaint is not a screen grab per repaint."""
now = fake_clock(modules, monkeypatch)
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, _entry = make_entity(modules, poller)
first = await entity.async_image()
now[0] += modules["image"].CAPTURE_REUSE_SECONDS - 1.0
second = await entity.async_image()
assert first == second == PNG_BYTES
assert len(poller.calls) == 1 # one shell, two image requests
async def test_the_window_expires_so_a_later_view_refreshes(self, modules, monkeypatch):
now = fake_clock(modules, monkeypatch)
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, _entry = make_entity(modules, poller)
await entity.async_image()
now[0] += modules["image"].CAPTURE_REUSE_SECONDS + 1.0
await entity.async_image()
assert len(poller.calls) == 2
async def test_an_explicit_refresh_skips_the_reuse_window(self, modules, monkeypatch):
"""``homeassistant.update_entity`` means "grab it now"."""
fake_clock(modules, monkeypatch)
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, _entry = make_entity(modules, poller)
await entity.async_image()
await entity.async_update()
assert len(poller.calls) == 2
async def test_concurrent_requests_share_one_capture(self, modules):
"""Parallel image fetches must not leak ADB shells."""
poller = FakeAdbPoller(wrapped(PNG_B64), delay=0.01)
entity, _entry = make_entity(modules, poller)
results = await asyncio.gather(
entity.async_image(), entity.async_image(), entity.async_image()
)
assert list(results) == [PNG_BYTES, PNG_BYTES, PNG_BYTES]
assert len(poller.calls) == 1
async def test_capture_reads_the_adb_poller_and_drives_nothing(self, modules):
"""The strict manager raises if the capture path touches anything else."""
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, _entry = make_entity(modules, poller) # StrictManager
assert await entity.async_image() == PNG_BYTES
# ...and the state HA would record is the timestamp, never the bytes
assert entity.image_last_updated is not None
assert isinstance(entity.image_last_updated, datetime)
assert entity.image_last_updated.tzinfo is not None
async def test_image_last_updated_is_stamped_only_by_a_success(self, modules, monkeypatch):
fake_clock(modules, monkeypatch)
poller = FakeAdbPoller(None)
entity, _entry = make_entity(modules, poller)
assert entity.image_last_updated is None # nothing captured yet
assert await entity.async_image() is None
assert entity.image_last_updated is None # a failure stamps nothing
poller.answer = wrapped(PNG_B64)
await entity.async_image()
stamped = entity.image_last_updated
assert stamped is not None
poller.answer = None # a later failure must not move it
await entity.async_image()
assert entity.image_last_updated is stamped
def test_entity_hygiene(self, modules):
"""Naming, content type, recorder hygiene, and the no-poll rule."""
poller = FakeAdbPoller(wrapped(PNG_B64))
entity, entry = make_entity(modules, poller)
cls = type(entity)
assert cls._attr_has_entity_name is True
assert entity.translation_key == "screen_capture"
assert entity.content_type == "image/png"
# HA polls image entities never by default; this one must not either
assert cls._attr_should_poll is False
assert entity.should_poll is False
assert entity.unique_id == f"nvidia_shield_{entry.entry_id}_screen_capture"
assert entity.device_info["identifiers"] == {("nvidia_shield", entry.entry_id)}
# the status changes with a capture, so it stays out of the database, and
# no attribute ever carries the image bytes
assert "capture_status" in cls._unrecorded_attributes
for value in entity.extra_state_attributes.values():
assert not isinstance(value, (bytes, bytearray))
def test_the_platform_is_forwarded(self, modules):
"""The new platform is in the entry's list, and the manifest still parses."""
platforms = [str(p) for p in sys.modules["nsh.init"].PLATFORMS]
assert "image" in platforms
assert platforms.count("image") == 1
manifest = json.loads(
(_SHIELD_DIR / "manifest.json").read_text(encoding="utf-8")
)
# HA discovers platforms from the forwarded list; a "platforms" key in the
# manifest is not a thing (and adding one would be a silent no-op).
assert "platforms" not in manifest
assert "adb-shell>=0.4.0" in manifest["requirements"]
# ── opt-in ───────────────────────────────────────────────────────────────────
class TestScreenCaptureOptIn:
"""Default OFF: no option, no entity, no cost."""
async def _setup(self, modules, options):
added: list = []
entry = make_entry(options)
entry.runtime_data = StrictManager(FakeAdbPoller(wrapped(PNG_B64)))
await modules["image"].async_setup_entry(SimpleNamespace(), entry, added.extend)
return added
async def test_the_option_defaults_off_and_creates_no_entity(self, modules):
assert await self._setup(modules, {}) == []
assert sys.modules["nsh.const"].CONF_SCREEN_CAPTURE_ENABLED == (
"screen_capture_enabled"
)
assert await self._setup(modules, {"screen_capture_enabled": False}) == []
async def test_switching_it_on_creates_exactly_one_entity(self, modules):
added = await self._setup(modules, {"screen_capture_enabled": True})
assert len(added) == 1
assert type(added[0]).__name__ == "NvidiaShieldScreenImage"
async def test_an_entry_without_a_manager_is_skipped_not_crashed(self, modules):
added: list = []
entry = make_entry({"screen_capture_enabled": True})
entry.runtime_data = None
await modules["image"].async_setup_entry(SimpleNamespace(), entry, added.extend)
assert added == []