feat(host-agent): synchronize device snapshots

2026-07-12 18:48:17 +08:00
parent f6bc8f6b98
commit 6e59f2f3ca
3 changed files with 136 additions and 1 deletions
@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from cloud.internal_api.models import DeviceSnapshotModel, HeartbeatResponse
+from device.manager import DeviceManager
+from host_agent.client import HostAgentClient
+from host_agent.config import HostAgentConfig
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+
+def build_device_snapshot(manager: DeviceManager) -> list[DeviceSnapshotModel]:
+    return [
+        DeviceSnapshotModel(
+            device_id=device.id,
+            driver_type=device.driver_type,
+            status=device.status,
+            capability_tags=list(device.capability_tags),
+        )
+        for device in sorted(manager.list_devices(), key=lambda item: item.id)
+    ]
+
+
+class HeartbeatSynchronizer:
+    def __init__(
+        self,
+        manager: DeviceManager,
+        client: HostAgentClient,
+        config: HostAgentConfig,
+        *,
+        address: str | None = None,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    ) -> None:
+        self.manager = manager
+        self.client = client
+        self.config = config
+        self.address = address
+        self._sleep = sleep
+
+    async def sync_once(self) -> HeartbeatResponse:
+        return await self.client.heartbeat(
+            build_device_snapshot(self.manager),
+            address=self.address,
+        )
+
+    async def run(self, stop: asyncio.Event) -> None:
+        while not stop.is_set():
+            await self.sync_once()
+            try:
+                await asyncio.wait_for(
+                    stop.wait(),
+                    timeout=self.config.heartbeat_interval_seconds,
+                )
+            except TimeoutError:
+                continue
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import UTC, datetime
+
+from cloud.internal_api.models import HeartbeatResponse
+from device.manager import DeviceManager
+from host_agent.config import HostAgentConfig
+from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot
+
+
+def _config() -> HostAgentConfig:
+    return HostAgentConfig(
+        control_plane_url="https://control.example",
+        host_id="host-a",
+        token="secret",
+        heartbeat_interval_seconds=0.01,
+    )
+
+
+def test_build_device_snapshot_copies_complete_non_secret_state() -> None:
+    manager = DeviceManager()
+    manager.register_device(
+        "device-b",
+        lambda: object(),  # type: ignore[arg-type,return-value]
+        driver_type="appium",
+        connection_info={"password": "must-not-leave-host"},
+        status="busy",
+        capability_tags=["android", "physical"],
+    )
+    manager.register_device(
+        "device-a",
+        lambda: object(),  # type: ignore[arg-type,return-value]
+        status="offline",
+    )
+
+    snapshot = build_device_snapshot(manager)
+
+    assert [device.device_id for device in snapshot] == ["device-a", "device-b"]
+    assert snapshot[1].driver_type == "appium"
+    assert snapshot[1].status == "busy"
+    assert snapshot[1].capability_tags == ["android", "physical"]
+    assert "must-not-leave-host" not in repr(snapshot)
+
+
+def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> None:
+    manager = DeviceManager()
+    manager.register_device(
+        "device-a",
+        lambda: object(),  # type: ignore[arg-type,return-value]
+    )
+    calls: list[list[str]] = []
+
+    class FakeClient:
+        async def heartbeat(self, devices, *, address=None):
+            calls.append([device.device_id for device in devices])
+            return HeartbeatResponse(
+                host_id="host-a",
+                accepted_devices=len(devices),
+                received_at=datetime.now(UTC),
+            )
+
+    async def scenario() -> None:
+        stop = asyncio.Event()
+        synchronizer = HeartbeatSynchronizer(
+            manager,
+            FakeClient(),  # type: ignore[arg-type]
+            _config(),
+        )
+        task = asyncio.create_task(synchronizer.run(stop))
+        while len(calls) < 3:
+            await asyncio.sleep(0.005)
+        stop.set()
+        await task
+
+    asyncio.run(scenario())
+    assert calls == [["device-a"], ["device-a"], ["device-a"]]