test(host-agent): cover distributed recovery flows
This commit is contained in:
@@ -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())
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user