fix: clean failed UIAutomator2 startups before retrying

tests/unit/clients/test_ui_automator_startup.py

@@ -0,0 +1,250 @@
+"""UIAutomator startup recovery, using the real library with isolated device IO."""
+
+from dataclasses import dataclass, field
+import json
+import threading
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+from uiautomator2 import base, core
+
+from artemis.clients import ui_automator_client
+from artemis.clients.accessibility_client import HelperUnavailable
+from artemis.clients.screen_client_factory import FallbackScreenClient
+
+
+@dataclass
+class Clock:
+    now: float = 0.0
+
+    def time(self):
+        return self.now
+
+    def sleep(self, seconds):
+        self.now += seconds
+
+
+class DeviceSocket:
+    def __init__(self, mode):
+        self.mode = mode
+        self.closed = threading.Event()
+        self.started = threading.Event()
+
+    def recv(self, _size):
+        if not self.started.is_set():
+            self.started.set()
+            output = b"[server] INFO: [UiAutomator2Server] Starting Server\n"
+            if self.mode == "conflict":
+                output += (
+                    b"java.lang.IllegalStateException: UiAutomationService already registered!"
+                )
+            return output
+        self.closed.wait()
+        return b""
+
+
+class DeviceConnection:
+    def __init__(self, mode):
+        self.conn = DeviceSocket(mode)
+        self.close_calls = 0
+        self.close_error = None
+
+    def close(self):
+        self.close_calls += 1
+        if self.close_error:
+            raise self.close_error
+        self.conn.closed.set()
+
+
+@dataclass
+class DeviceServer:
+    modes: list[str] = field(default_factory=lambda: ["not_ready"])
+    connections: list[DeviceConnection] = field(default_factory=list)
+    previous_closed_at_launch: list[bool] = field(default_factory=list)
+    clock: Clock = field(default_factory=Clock)
+    cleanup_callbacks: MagicMock = field(default_factory=MagicMock)
+    close_error: Exception | None = None
+    probe_error: BaseException | None = None
+
+    def shell(self, command, stream=False):
+        assert stream is True
+        assert "com.wetest.uia2.Main" in command
+        self.previous_closed_at_launch.append(all(c.conn.closed.is_set() for c in self.connections))
+        index = min(len(self.connections), len(self.modes) - 1)
+        connection = DeviceConnection(self.modes[index])
+        connection.close_error = self.close_error
+        self.connections.append(connection)
+        return connection
+
+    def http_request(self, _device, _port, method, path, *args, **kwargs):
+        if self.connections:
+            latest = self.connections[-1]
+            assert latest.conn.started.wait(timeout=2)
+            if self.probe_error is not None:
+                raise self.probe_error
+            if latest.conn.mode == "ready" and not latest.conn.closed.is_set():
+                if method == "POST":
+                    return core.HTTPResponse(
+                        json.dumps({"result": {"displayWidth": 1080}}).encode()
+                    )
+                return core.HTTPResponse(b"pong")
+        raise core.HTTPError("simulated device server is not ready")
+
+
+@pytest.fixture
+def server(monkeypatch):
+    simulated = DeviceServer()
+    device = SimpleNamespace(serial="isolated-startup-device", shell=simulated.shell)
+    monkeypatch.setattr(base._BaseClient, "_wait_for_device", lambda _self: device)
+    monkeypatch.setattr(core.BasicUiautomatorServer, "_setup_jar", lambda _self: None)
+    monkeypatch.setattr(core, "_http_request", simulated.http_request)
+    monkeypatch.setattr(core, "time", simulated.clock)
+    monkeypatch.setattr(core, "atexit", SimpleNamespace(register=simulated.cleanup_callbacks))
+    monkeypatch.setattr(ui_automator_client, "time", simulated.clock)
+    monkeypatch.setattr(ui_automator_client, "_ensure_maestro_not_installed", lambda _s: None)
+    monkeypatch.setattr(ui_automator_client, "ensure_device_awake", lambda _s: "simulated")
+    yield simulated
+    # Close fake sockets even when a regression leaves a real library reader alive.
+    for connection in simulated.connections:
+        connection.conn.closed.set()
+
+
+def test_each_failed_launch_is_closed_before_retry_and_final_failure(server):
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    with pytest.raises(RuntimeError, match="server not ready") as caught:
+        client.connect()
+
+    assert len(server.connections) == 3
+    assert server.previous_closed_at_launch == [True, True, True]
+    assert [c.close_calls for c in server.connections] == [1, 1, 1]
+    assert server.clock.now == 93.0
+    assert server.cleanup_callbacks.call_count == 0
+    assert client._device is None
+    assert client._awake_strategy is None
+    assert isinstance(caught.value.__cause__, core.LaunchUiAutomationError)
+    assert "isolated-startup-device" in str(caught.value)
+    assert "3" in str(caught.value)
+
+
+def test_retry_can_connect_after_failed_launch_was_stopped(server):
+    server.modes = ["not_ready", "ready"]
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    client.connect()
+    client.connect()  # Cached device info also uses the real library RPC path.
+
+    assert len(server.connections) == 2
+    assert server.previous_closed_at_launch == [True, True]
+    assert server.connections[0].close_calls == 1
+    assert server.connections[1].close_calls == 0
+    assert server.cleanup_callbacks.call_count == 1
+    client.disconnect(stop_server=True)
+    assert server.connections[1].close_calls == 1
+
+
+def test_later_connection_can_succeed_after_all_previous_attempts_failed(server):
+    server.modes = ["not_ready", "not_ready", "not_ready", "ready"]
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    with pytest.raises(RuntimeError, match="server not ready"):
+        client.connect()
+    client.connect()
+
+    assert len(server.connections) == 4
+    assert server.previous_closed_at_launch == [True, True, True, True]
+    assert [c.close_calls for c in server.connections] == [1, 1, 1, 0]
+    client.disconnect(stop_server=True)
+    assert server.connections[-1].close_calls == 1
+
+
+def test_auto_fallback_cleans_failed_launches_before_next_attempt(server):
+    helper = MagicMock()
+    helper.connect.side_effect = HelperUnavailable("simulated helper unavailable")
+    client = FallbackScreenClient(
+        "isolated-startup-device",
+        helper=helper,
+        uiautomator_factory=lambda: ui_automator_client.UIAutomatorClient(
+            "isolated-startup-device"
+        ),
+        state_probe=lambda _serial: "device",
+    )
+
+    with pytest.raises(RuntimeError, match="server not ready"):
+        client.connect()
+    client.disconnect()
+
+    assert len(server.connections) == 3
+    assert server.previous_closed_at_launch == [True, True, True]
+    assert [c.close_calls for c in server.connections] == [1, 1, 1]
+    helper.disconnect.assert_called_once()
+
+
+def test_service_conflict_closes_own_streams_and_preserves_cause(server):
+    server.modes = ["conflict"]
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    with pytest.raises(RuntimeError, match="already registered") as caught:
+        client.connect()
+
+    assert isinstance(caught.value.__cause__, core.AccessibilityServiceAlreadyRegisteredError)
+    assert len(server.connections) == 3
+    assert [c.close_calls for c in server.connections] == [1, 1, 1]
+
+
+def test_cleanup_error_does_not_replace_original_startup_failure(server):
+    server.close_error = OSError("simulated cleanup connection error")
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    with pytest.raises(RuntimeError, match="server not ready") as caught:
+        client.connect()
+
+    assert isinstance(caught.value.__cause__, core.LaunchUiAutomationError)
+    assert [c.close_calls for c in server.connections] == [1, 1, 1]
+
+
+def test_interrupt_during_startup_also_closes_own_stream(server):
+    server.probe_error = KeyboardInterrupt()
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    with pytest.raises(KeyboardInterrupt):
+        client.connect()
+
+    assert len(server.connections) == 1
+    assert server.connections[0].close_calls == 1
+
+
+def test_failure_before_launch_does_not_stop_another_device_service(server, monkeypatch):
+    failure = OSError("simulated jar transfer failure")
+
+    def fail_before_launch(_self):
+        raise failure
+
+    monkeypatch.setattr(core.BasicUiautomatorServer, "_setup_jar", fail_before_launch)
+    stop = MagicMock()
+    monkeypatch.setattr(core.BasicUiautomatorServer, "stop_uiautomator", stop)
+    client = ui_automator_client.UIAutomatorClient("isolated-startup-device")
+
+    with pytest.raises(RuntimeError, match="jar transfer failure") as caught:
+        client.connect()
+
+    assert caught.value.__cause__ is failure
+    assert server.connections == []
+    stop.assert_not_called()
+
+
+def test_default_helper_success_does_not_start_uiautomator(server):
+    factory = MagicMock()
+    client = FallbackScreenClient(
+        "isolated-startup-device",
+        helper=MagicMock(),
+        uiautomator_factory=factory,
+        state_probe=lambda _serial: "device",
+    )
+
+    client.connect()
+    client.disconnect()
+
+    factory.assert_not_called()
+    assert server.connections == []

tests/unit/test_ui_automator_client_awake.py

@@ -6,7 +6,7 @@ from artemis.clients.ui_automator_client import UIAutomatorClient
 from artemis.drivers.android.adb_driver import AndroidAdbDriver
 
 
-@patch("artemis.clients.ui_automator_client.u2.connect")
+@patch("artemis.clients.ui_automator_client._ManagedUIAutomatorDevice")
 @patch(
     "artemis.clients.ui_automator_client.ensure_device_awake",
     return_value="host_heartbeat",
@@ -25,7 +25,10 @@ def test_new_ui_connection_enrolls_device_in_shared_awake_strategy(
     mock_connect.assert_called_once_with("device-123")
 
 
-@patch("artemis.clients.ui_automator_client.u2.connect", side_effect=RuntimeError("offline"))
+@patch(
+    "artemis.clients.ui_automator_client._ManagedUIAutomatorDevice",
+    side_effect=RuntimeError("offline"),
+)
 @patch(
     "artemis.clients.ui_automator_client.ensure_device_awake",
     return_value="host_heartbeat",
@@ -43,7 +46,7 @@ def test_failed_ui_connection_does_not_stop_process_awake_service(
     assert client._awake_strategy is None
 
 
-@patch("artemis.clients.ui_automator_client.u2.connect")
+@patch("artemis.clients.ui_automator_client._ManagedUIAutomatorDevice")
 @patch(
     "artemis.clients.ui_automator_client.ensure_device_awake",
     return_value="host_heartbeat",