Implement DroidCast Screenshot Backend

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")

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")