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,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"]]