326 lines
11 KiB
Python
326 lines
11 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 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())
|