feat(host-agent): add control plane client

2026-07-12 18:44:18 +08:00
parent 1e10d5aa95
commit f6bc8f6b98
4 changed files with 335 additions and 5 deletions
@@ -0,0 +1,133 @@
+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