466 lines
16 KiB
Python
466 lines
16 KiB
Python
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 fastapi.testclient import TestClient
|
|
|
|
from cloud.auth import BearerCredential
|
|
from cloud.control_config import CloudControlConfig
|
|
from cloud.sdk.client import CloudClient
|
|
from cloud.scheduler import TaskConstraints
|
|
from cloud_api.app import create_app
|
|
from core.models import Scene
|
|
from device.manager import DeviceManager
|
|
from driver.base import Driver
|
|
from host_agent.assignment import AssignmentExecutionResult, AssignmentExecutor
|
|
from host_agent.client import HostAgentClient, StaleLeaseError
|
|
from host_agent.config import HostAgentConfig
|
|
from host_agent.execution import ExecutionFactories
|
|
from host_agent.heartbeat import HeartbeatSynchronizer
|
|
from host_agent.lease import ActiveAssignmentRunner
|
|
from host_agent.processor import AssignmentProcessor
|
|
from runtime.executor import Executor, default_tool_registry
|
|
from runtime.planner import PlannedStep, Planner
|
|
from runtime.task import TaskRunner, TaskRunnerConfig
|
|
from workflow.store import WorkflowStore
|
|
|
|
|
|
class FakeDriver(Driver):
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, tuple[object, ...]]] = []
|
|
|
|
def connect(self) -> None:
|
|
self.calls.append(("connect", ()))
|
|
|
|
def disconnect(self) -> None:
|
|
return None
|
|
|
|
def screenshot(self) -> bytes:
|
|
return b""
|
|
|
|
def tap(self, x: float, y: float) -> None:
|
|
self.calls.append(("tap", (x, y)))
|
|
|
|
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 _public_credential() -> BearerCredential:
|
|
return BearerCredential(
|
|
principal_id="sdk",
|
|
token="public-token",
|
|
scopes=frozenset({"tasks:submit", "tasks:read"}),
|
|
)
|
|
|
|
|
|
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
|
|
|
|
|
|
class _AsyncTestClient:
|
|
def __init__(self, client: TestClient) -> None:
|
|
self.client = client
|
|
|
|
async def request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
|
kwargs.pop("timeout", None)
|
|
return await asyncio.to_thread(self.client.request, method, path, **kwargs)
|
|
|
|
|
|
class _TapPlanner(Planner):
|
|
def plan(self, *, goal, scene, context):
|
|
if context.step_results:
|
|
return []
|
|
return [PlannedStep(action="tap", description="tap", args={"x": 2, "y": 3})]
|
|
|
|
def goal_reached(self, *, goal, scene, context):
|
|
return bool(context.step_results)
|
|
|
|
|
|
class _FailingPlanner(Planner):
|
|
def plan(self, *, goal, scene, context):
|
|
raise RuntimeError("planner unavailable")
|
|
|
|
|
|
def _assignment_processor(
|
|
tmp_path: Path,
|
|
client: HostAgentClient,
|
|
manager: DeviceManager,
|
|
*,
|
|
planner: Planner,
|
|
) -> AssignmentProcessor:
|
|
scene = Scene(width=10, height=20, elements=[])
|
|
|
|
def task_runner_factory() -> TaskRunner:
|
|
return TaskRunner(
|
|
planner=planner,
|
|
executor=Executor(tools=default_tool_registry(manager=manager)),
|
|
observer=lambda device_id: scene,
|
|
screenshot_provider=lambda device_id: b"",
|
|
config=TaskRunnerConfig(max_steps=3),
|
|
)
|
|
|
|
workflow_store = WorkflowStore(tmp_path / "workflows.sqlite3")
|
|
factories = ExecutionFactories(
|
|
task_runner_factory=task_runner_factory,
|
|
workflow_runner_factory=lambda: object(), # type: ignore[arg-type,return-value]
|
|
workflow_store=workflow_store,
|
|
)
|
|
executor = AssignmentExecutor(factories)
|
|
return AssignmentProcessor(client, ActiveAssignmentRunner(client, executor))
|
|
|
|
|
|
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())
|
|
|
|
|
|
def test_public_sdk_reports_fake_device_success_and_runtime_failure(tmp_path) -> None:
|
|
app = create_app(
|
|
config=CloudControlConfig(
|
|
environment="test",
|
|
database_url=f"sqlite:///{(tmp_path / 'sdk-e2e.sqlite3').as_posix()}",
|
|
scheduler_interval_seconds=60,
|
|
lease_reaper_interval_seconds=60,
|
|
lease_duration_seconds=30,
|
|
credentials=(_public_credential(), _credential("host-a")),
|
|
)
|
|
)
|
|
driver = FakeDriver()
|
|
manager = DeviceManager()
|
|
manager.register_device("device-a", lambda: driver, status="idle")
|
|
|
|
with TestClient(app) as http_client:
|
|
cloud_client = CloudClient(
|
|
"http://testserver",
|
|
http_client=http_client,
|
|
token="public-token",
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
host_client = HostAgentClient(
|
|
_config("host-a"),
|
|
http_client=_AsyncTestClient(http_client), # type: ignore[arg-type]
|
|
)
|
|
heartbeat = HeartbeatSynchronizer(manager, host_client, _config("host-a"))
|
|
heartbeat.connect_devices()
|
|
await heartbeat.sync_once()
|
|
|
|
successful_task_id = cloud_client.submit_task(goal="tap screen")[
|
|
"task_id"
|
|
]
|
|
app.state.cloud_services.scheduler.assign()
|
|
successful_assignment = await host_client.claim()
|
|
assert successful_assignment is not None
|
|
await _assignment_processor(
|
|
tmp_path / "success",
|
|
host_client,
|
|
manager,
|
|
planner=_TapPlanner(),
|
|
).process(successful_assignment)
|
|
|
|
successful_status = cloud_client.get_task_status(successful_task_id)
|
|
assert successful_status["status"] == "done"
|
|
assert ("tap", (2, 3)) in driver.calls
|
|
|
|
failed_task_id = cloud_client.submit_task(goal="fail planning")["task_id"]
|
|
app.state.cloud_services.scheduler.assign()
|
|
failed_assignment = await host_client.claim()
|
|
assert failed_assignment is not None
|
|
await _assignment_processor(
|
|
tmp_path / "failure",
|
|
host_client,
|
|
manager,
|
|
planner=_FailingPlanner(),
|
|
).process(failed_assignment)
|
|
|
|
failed_status = cloud_client.get_task_status(failed_task_id)
|
|
assert failed_status["status"] == "failed"
|
|
assert failed_status["failure_reason"] == (
|
|
"RuntimeError: planner unavailable"
|
|
)
|
|
|
|
asyncio.run(scenario())
|