diff --git a/apps/device-host-agent/tests/test_e2e.py b/apps/device-host-agent/tests/test_e2e.py new file mode 100644 index 0000000..63c53a8 --- /dev/null +++ b/apps/device-host-agent/tests/test_e2e.py @@ -0,0 +1,321 @@ +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import UTC, datetime +from pathlib import Path + +import httpx +import pytest + +from cloud.auth import BearerCredential +from cloud.control_config import CloudControlConfig +from cloud.scheduler import TaskConstraints +from cloud_api.app import create_app +from device.manager import DeviceManager +from driver.base import Driver +from host_agent.assignment import AssignmentExecutionResult +from host_agent.client import HostAgentClient, StaleLeaseError +from host_agent.config import HostAgentConfig +from host_agent.heartbeat import HeartbeatSynchronizer +from host_agent.processor import AssignmentProcessor + + +class FakeDriver(Driver): + def connect(self) -> None: + return None + + def disconnect(self) -> None: + return None + + def screenshot(self) -> bytes: + return b"" + + def tap(self, x: float, y: float) -> None: + return None + + def swipe( + self, + start_x: float, + start_y: float, + end_x: float, + end_y: float, + duration_ms: int = 500, + ) -> None: + return None + + def input(self, text: str) -> None: + return None + + def launch(self, app_id: str) -> None: + return None + + def terminate(self, app_id: str) -> None: + return None + + def tree(self): + return None + + def home(self) -> None: + return None + + def lock(self) -> None: + return None + + def unlock(self) -> None: + return None + + +def _credential(host_id: str) -> BearerCredential: + return BearerCredential( + principal_id=f"agent-{host_id}", + token=f"token-{host_id}", + scopes=frozenset(), + host_id=host_id, + ) + + +def _config(host_id: str) -> HostAgentConfig: + return HostAgentConfig( + control_plane_url="http://control.test", + host_id=host_id, + token=f"token-{host_id}", + poll_timeout_seconds=0.01, + retry_backoff_seconds=0.001, + max_retry_backoff_seconds=0.001, + ) + + +@asynccontextmanager +async def _control_plane( + database_path: Path, + *host_ids: str, + lease_duration_seconds: float = 30, +) -> AsyncIterator[object]: + app = create_app( + config=CloudControlConfig( + environment="test", + database_url=f"sqlite:///{database_path.as_posix()}", + scheduler_interval_seconds=60, + lease_reaper_interval_seconds=60, + lease_duration_seconds=lease_duration_seconds, + credentials=tuple(_credential(host_id) for host_id in host_ids), + ) + ) + async with app.router.lifespan_context(app): + yield app + + +@asynccontextmanager +async def _host_client( + app, + host_id: str, + *, + request_paths: list[str] | None = None, +) -> AsyncIterator[HostAgentClient]: + transport: httpx.AsyncBaseTransport = httpx.ASGITransport(app=app) + if request_paths is not None: + transport = _RecordingTransport(transport, request_paths) + async with httpx.AsyncClient( + transport=transport, + base_url="http://control.test", + ) as http_client: + yield HostAgentClient(_config(host_id), http_client=http_client) + + +class _RecordingTransport(httpx.AsyncBaseTransport): + def __init__( + self, + transport: httpx.AsyncBaseTransport, + paths: list[str], + ) -> None: + self.transport = transport + self.paths = paths + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + self.paths.append(request.url.path) + return await self.transport.handle_async_request(request) + + async def aclose(self) -> None: + await self.transport.aclose() + + +async def _sync_fake_device( + client: HostAgentClient, + host_id: str, + device_id: str, + *, + driver_type: str = "wda", +) -> FakeDriver: + driver = FakeDriver() + manager = DeviceManager() + manager.register_device( + device_id, + lambda: driver, + driver_type=driver_type, + status="idle", + ) + await HeartbeatSynchronizer(manager, client, _config(host_id)).sync_once() + return driver + + +class _SuccessfulExecution: + async def run(self, assignment): + return AssignmentExecutionResult( + status="done", + metadata={"fake_driver": assignment.device_id}, + ) + + def request_stop(self) -> None: + return None + + +def test_one_host_executes_assignment_through_outbound_protocol(tmp_path) -> None: + async def scenario() -> None: + paths: list[str] = [] + async with _control_plane(tmp_path / "one-host.sqlite3", "host-a") as app: + async with _host_client(app, "host-a", request_paths=paths) as client: + await _sync_fake_device(client, "host-a", "device-a") + task_id = app.state.cloud_services.scheduler.submit(goal="open settings") + app.state.cloud_services.scheduler.assign() + assignment = await client.claim() + assert assignment is not None + + result = await AssignmentProcessor( + client, + _SuccessfulExecution(), + ).process(assignment) + + task = app.state.cloud_services.repository.get_task(task_id) + assert result.report_status == "recorded" + assert task.status == "done" + assert task.terminal_result == {"fake_driver": "device-a"} + assert all(path.startswith("/internal/v1/") for path in paths) + + asyncio.run(scenario()) + + +def test_nat_style_host_requires_only_outbound_requests(tmp_path) -> None: + async def scenario() -> None: + paths: list[str] = [] + async with _control_plane(tmp_path / "outbound-only.sqlite3", "host-a") as app: + async with _host_client(app, "host-a", request_paths=paths) as client: + await _sync_fake_device(client, "host-a", "device-a") + assert await client.claim() is None + + assert paths == [ + "/internal/v1/hosts/host-a/heartbeat", + "/internal/v1/hosts/host-a/assignments/claim", + ] + + asyncio.run(scenario()) + + +def test_multiple_hosts_claim_only_their_matching_devices(tmp_path) -> None: + async def scenario() -> None: + async with _control_plane( + tmp_path / "multiple-hosts.sqlite3", + "host-a", + "host-b", + ) as app: + async with ( + _host_client(app, "host-a") as client_a, + _host_client(app, "host-b") as client_b, + ): + await _sync_fake_device(client_a, "host-a", "device-a") + await _sync_fake_device( + client_b, + "host-b", + "device-b", + driver_type="appium", + ) + task_a = app.state.cloud_services.scheduler.submit( + goal="ios task", + constraints=TaskConstraints(driver_type="wda"), + ) + task_b = app.state.cloud_services.scheduler.submit( + goal="android task", + constraints=TaskConstraints(driver_type="appium"), + ) + app.state.cloud_services.scheduler.assign() + + assignment_a = await client_a.claim() + assignment_b = await client_b.claim() + + assert assignment_a is not None and assignment_a.task_id == task_a + assert assignment_b is not None and assignment_b.task_id == task_b + assert assignment_a.host_id == "host-a" + assert assignment_b.host_id == "host-b" + + asyncio.run(scenario()) + + +def test_control_plane_restart_preserves_dispatched_assignment(tmp_path) -> None: + async def scenario() -> None: + database_path = tmp_path / "control-restart.sqlite3" + assignment = None + async with _control_plane(database_path, "host-a") as first_app: + async with _host_client(first_app, "host-a") as client: + await _sync_fake_device(client, "host-a", "device-a") + task_id = first_app.state.cloud_services.scheduler.submit(goal="resume") + first_app.state.cloud_services.scheduler.assign() + assignment = await client.claim() + assert assignment is not None + + async with _control_plane(database_path, "host-a") as restarted_app: + async with _host_client(restarted_app, "host-a") as client: + assert (await client.renew(assignment)).status == "renewed" + await client.report_result(assignment, status="done") + task = restarted_app.state.cloud_services.repository.get_task(task_id) + assert task.status == "done" + + asyncio.run(scenario()) + + +def test_host_agent_restart_reuses_active_lease(tmp_path) -> None: + async def scenario() -> None: + async with _control_plane(tmp_path / "agent-restart.sqlite3", "host-a") as app: + async with _host_client(app, "host-a") as first_client: + await _sync_fake_device(first_client, "host-a", "device-a") + task_id = app.state.cloud_services.scheduler.submit(goal="resume host") + app.state.cloud_services.scheduler.assign() + assignment = await first_client.claim() + assert assignment is not None + + async with _host_client(app, "host-a") as restarted_client: + assert (await restarted_client.renew(assignment)).status == "renewed" + await restarted_client.report_result(assignment, status="done") + + assert app.state.cloud_services.repository.get_task(task_id).status == "done" + + asyncio.run(scenario()) + + +def test_lease_loss_rejects_stale_host_result(tmp_path) -> None: + async def scenario() -> None: + async with _control_plane( + tmp_path / "lease-loss.sqlite3", + "host-a", + lease_duration_seconds=0.2, + ) as app: + async with _host_client(app, "host-a") as client: + await _sync_fake_device(client, "host-a", "device-a") + task_id = app.state.cloud_services.scheduler.submit(goal="expire") + app.state.cloud_services.scheduler.assign() + assignment = await client.claim() + assert assignment is not None + await asyncio.sleep(0.25) + app.state.cloud_services.repository.reap_expired_leases( + now=datetime.now(UTC), + max_attempts=3, + ) + + with pytest.raises(StaleLeaseError): + await client.report_result(assignment, status="done") + + task = app.state.cloud_services.repository.get_task(task_id) + assert task.status == "queued" + assert task.terminal_result is None + + asyncio.run(scenario()) diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index cfa680d..584b158 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -59,7 +59,7 @@ - [x] 7.5 Run lease renewal alongside active execution and stop further interruptible actions after confirmed lease loss. - [x] 7.6 Normalize and report successful/failed terminal outcomes, including Runtime failure reasons, with idempotent retries after response loss. - [x] 7.7 Implement graceful shutdown that stops polling, finishes or interrupts current work according to lease policy, and performs a final heartbeat when possible. -- [ ] 7.8 Add fake-driver end-to-end tests for one host, multiple hosts, NAT-style outbound-only operation, control-plane restart, Host Agent restart, and lease loss. +- [x] 7.8 Add fake-driver end-to-end tests for one host, multiple hosts, NAT-style outbound-only operation, control-plane restart, Host Agent restart, and lease loss. ## 8. Public SDK And Operational Delivery