78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
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"]]
|