This commit is contained in:
@@ -7,6 +7,7 @@ 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
|
||||
from host_agent.status import AgentStatusTracker
|
||||
|
||||
|
||||
class ConnectableDriver:
|
||||
@@ -81,3 +82,44 @@ def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> N
|
||||
asyncio.run(scenario())
|
||||
assert calls == [["device-a"], ["device-a"], ["device-a"]]
|
||||
assert manager.status("device-a") == "busy"
|
||||
|
||||
|
||||
def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> None:
|
||||
manager = DeviceManager()
|
||||
manager.register_device(
|
||||
"device-a",
|
||||
lambda: ConnectableDriver(), # type: ignore[arg-type,return-value]
|
||||
)
|
||||
manager.register_device(
|
||||
"device-b",
|
||||
lambda: ConnectableDriver(), # type: ignore[arg-type,return-value]
|
||||
)
|
||||
|
||||
class FakeClient:
|
||||
async def heartbeat(self, devices, *, address=None):
|
||||
return HeartbeatResponse(
|
||||
host_id="host-a",
|
||||
accepted_devices=len(devices),
|
||||
received_at=datetime.now(UTC),
|
||||
)
|
||||
|
||||
async def scenario() -> None:
|
||||
tracker = AgentStatusTracker()
|
||||
on_sync_calls: list[int] = []
|
||||
synchronizer = HeartbeatSynchronizer(
|
||||
manager,
|
||||
FakeClient(), # type: ignore[arg-type]
|
||||
_config(),
|
||||
status_tracker=tracker,
|
||||
on_sync=on_sync_calls.append,
|
||||
)
|
||||
|
||||
await synchronizer.sync_once()
|
||||
|
||||
assert on_sync_calls == [2]
|
||||
last_heartbeat = tracker.snapshot()["last_heartbeat"]
|
||||
assert last_heartbeat is not None
|
||||
assert last_heartbeat["ok"] is True
|
||||
assert last_heartbeat["device_count"] == 2
|
||||
|
||||
asyncio.run(scenario())
|
||||
|
||||
Reference in New Issue
Block a user