diff --git a/adb_vision/__init__.py b/adb_vision/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/adb_vision/screenshot.py b/adb_vision/screenshot.py new file mode 100644 index 0000000000..bcb2c89b28 --- /dev/null +++ b/adb_vision/screenshot.py @@ -0,0 +1,70 @@ +import asyncio +import base64 +import logging +import urllib.request +from typing import Callable, Awaitable + +# Type alias for the async ADB runner +AdbRunFn = Callable[..., Awaitable[bytes]] + +async def _capture_droidcast(*, adb_run: AdbRunFn, serial: str, adb_exe: str) -> str: + """Return base64-encoded PNG screenshot via DroidCast HTTP stream.""" + port = "53516" + apk_path = "alas_wrapped/bin/DroidCast/DroidCast-debug-1.2.3.apk" + + # 1. Idempotency check: if port forward already exists, try to fetch first + try: + forward_list = await adb_run("forward", "--list", timeout=5.0) + if f"tcp:{port}" in forward_list.decode(): + try: + return await _fetch_screenshot(port) + except Exception: + # Port forward exists but service might not be responding, continue to start + pass + except Exception: + pass + + # 2. Push/Install APK if needed + # (In a real scenario we'd check if installed, but following the plan's direct install -r) + try: + await adb_run("install", "-r", apk_path, timeout=30.0) + except Exception as e: + # If the file is missing in the environment, we log it but might proceed if it's already installed + logging.warning(f"DroidCast install failed: {e}") + + # 3. Start DroidCast service + await adb_run( + "shell", "am", "start-foreground-service", + "-n", "com.rayworks.droidcast/.Main", + "-a", "android.intent.action.MAIN", + timeout=10.0 + ) + + # 4. Forward the port + await adb_run("forward", f"tcp:{port}", f"tcp:{port}", timeout=5.0) + + # 5. Fetch and return + # Allow a small retry loop for the service to bind the port + for i in range(5): + try: + return await _fetch_screenshot(port) + except Exception: + if i == 4: + raise + await asyncio.sleep(0.5) + + +async def _fetch_screenshot(port: str) -> str: + """Internal helper to GET screenshot from localhost:port.""" + url = f"http://localhost:{port}/screenshot?format=png" + + def _get(): + # Using urllib.request to avoid extra dependencies like aiohttp + with urllib.request.urlopen(url, timeout=5.0) as resp: + if resp.status != 200: + raise RuntimeError(f"DroidCast HTTP error: {resp.status}") + return resp.read() + + loop = asyncio.get_running_loop() + data = await loop.run_in_executor(None, _get) + return base64.b64encode(data).decode("ascii") diff --git a/adb_vision/test_droidcast.py b/adb_vision/test_droidcast.py new file mode 100644 index 0000000000..a763beb40c --- /dev/null +++ b/adb_vision/test_droidcast.py @@ -0,0 +1,98 @@ +import asyncio +import base64 +import io +import pytest +from unittest import mock +from adb_vision.screenshot import _capture_droidcast + +def _fake_png() -> bytes: + """Return a minimal valid PNG (1x1 white square).""" + return ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00" + b"\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc\xf8" + b"\xff\xff\x3f\x00\x05\xfe\x02\xfe\x05\x1e\x35\xaf\x00\x00\x00\x00IEND\xaeB`\x82" + ) + +@pytest.mark.asyncio +async def test_capture_droidcast_happy_path(): + """Verify happy path: install -> start -> forward -> fetch.""" + adb_run = mock.AsyncMock() + adb_run.return_value = b"" # Default for all calls + + # Mocking --list to return empty initially + adb_run.side_effect = [ + b"", # forward --list + b"", # install + b"", # am start-foreground-service + b"", # forward tcp:53516 tcp:53516 + ] + + png_data = _fake_png() + + # Mock urllib.request.urlopen + with mock.patch("urllib.request.urlopen") as mock_url: + mock_resp = mock.MagicMock() + mock_resp.status = 200 + mock_resp.read.return_value = png_data + mock_resp.__enter__.return_value = mock_resp + mock_url.return_value = mock_resp + + result = await _capture_droidcast(adb_run=adb_run, serial="emulator-5554", adb_exe="adb") + + assert result == base64.b64encode(png_data).decode("ascii") + + # Verify calls + # Call 0: forward --list + # Call 1: install -r + # Call 2: am start-foreground-service + # Call 3: forward tcp:53516 tcp:53516 + assert adb_run.call_count == 4 + adb_run.assert_any_call("install", "-r", mock.ANY, timeout=30.0) + adb_run.assert_any_call("shell", "am", "start-foreground-service", "-n", "com.rayworks.droidcast/.Main", "-a", "android.intent.action.MAIN", timeout=10.0) + adb_run.assert_any_call("forward", "tcp:53516", "tcp:53516", timeout=5.0) + +@pytest.mark.asyncio +async def test_capture_droidcast_idempotent(): + """Verify it skips install/start if port is already forwarded and responsive.""" + adb_run = mock.AsyncMock() + adb_run.return_value = b"tcp:53516 tcp:53516" # forward --list says it exists + + png_data = _fake_png() + + with mock.patch("urllib.request.urlopen") as mock_url: + mock_resp = mock.MagicMock() + mock_resp.status = 200 + mock_resp.read.return_value = png_data + mock_resp.__enter__.return_value = mock_resp + mock_url.return_value = mock_resp + + result = await _capture_droidcast(adb_run=adb_run, serial="emulator-5554", adb_exe="adb") + + assert result == base64.b64encode(png_data).decode("ascii") + + # Should only have called forward --list + assert adb_run.call_count == 1 + adb_run.assert_called_once_with("forward", "--list", timeout=5.0) + +@pytest.mark.asyncio +async def test_capture_droidcast_http_failure(): + """Verify it raises RuntimeError on HTTP failure.""" + adb_run = mock.AsyncMock() + adb_run.return_value = b"" + + # Mocking --list to return empty initially + adb_run.side_effect = [ + b"", # forward --list + b"", # install + b"", # am start-foreground-service + b"", # forward tcp:53516 tcp:53516 + ] + + with mock.patch("urllib.request.urlopen") as mock_url: + mock_resp = mock.MagicMock() + mock_resp.status = 404 # Failure + mock_resp.__enter__.return_value = mock_resp + mock_url.return_value = mock_resp + + with pytest.raises(RuntimeError, match="DroidCast HTTP error: 404"): + await _capture_droidcast(adb_run=adb_run, serial="emulator-5554", adb_exe="adb")