134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import UTC, datetime
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from cloud.internal_api.models import AssignmentModel, DeviceSnapshotModel
|
|
from host_agent.client import HostAgentClient, StaleLeaseError
|
|
from host_agent.config import HostAgentConfig
|
|
|
|
|
|
def _config(**overrides) -> HostAgentConfig:
|
|
values = {
|
|
"control_plane_url": "https://control.example",
|
|
"host_id": "host-a",
|
|
"token": "host-secret",
|
|
"retry_backoff_seconds": 0.01,
|
|
"max_retry_backoff_seconds": 0.02,
|
|
"max_retry_attempts": 3,
|
|
}
|
|
values.update(overrides)
|
|
return HostAgentConfig(**values)
|
|
|
|
|
|
def _assignment() -> AssignmentModel:
|
|
return AssignmentModel(
|
|
task_id="task-a",
|
|
attempt=1,
|
|
lease_id="lease-a",
|
|
lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC),
|
|
host_id="host-a",
|
|
device_id="device-a",
|
|
goal="open settings",
|
|
)
|
|
|
|
|
|
def test_client_sends_authenticated_heartbeat_and_claim() -> None:
|
|
requests: list[httpx.Request] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
if request.url.path.endswith("/heartbeat"):
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"host_id": "host-a",
|
|
"accepted_devices": 1,
|
|
"received_at": "2026-07-12T00:00:00Z",
|
|
},
|
|
)
|
|
return httpx.Response(200, json={"assignment": None, "timed_out": True})
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
heartbeat = await client.heartbeat(
|
|
[
|
|
DeviceSnapshotModel(
|
|
device_id="device-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
)
|
|
]
|
|
)
|
|
assignment = await client.claim()
|
|
assert heartbeat.accepted_devices == 1
|
|
assert assignment is None
|
|
|
|
asyncio.run(scenario())
|
|
assert [request.url.path for request in requests] == [
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
"/internal/v1/hosts/host-a/assignments/claim",
|
|
]
|
|
assert all(
|
|
request.headers["authorization"] == "Bearer host-secret" for request in requests
|
|
)
|
|
|
|
|
|
def test_transport_failures_retry_with_bounded_backoff() -> None:
|
|
attempts = 0
|
|
sleeps: list[float] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts < 3:
|
|
raise httpx.ConnectError("temporary", request=request)
|
|
return httpx.Response(200, json={"assignment": None, "timed_out": True})
|
|
|
|
async def sleep(delay: float) -> None:
|
|
sleeps.append(delay)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client, sleep=sleep)
|
|
assert await client.claim() is None
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 3
|
|
assert sleeps == [0.01, 0.02]
|
|
|
|
|
|
def test_stale_lease_response_raises_typed_error_without_retry() -> None:
|
|
attempts = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
return httpx.Response(
|
|
409,
|
|
json={"code": "stale_lease", "detail": "lease expired"},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
with pytest.raises(StaleLeaseError) as error:
|
|
await client.renew(_assignment())
|
|
assert "host-secret" not in str(error.value)
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|