diff --git a/apps/device-host-agent/host_agent/heartbeat.py b/apps/device-host-agent/host_agent/heartbeat.py index a8dec71..d3c249c 100644 --- a/apps/device-host-agent/host_agent/heartbeat.py +++ b/apps/device-host-agent/host_agent/heartbeat.py @@ -4,6 +4,7 @@ import asyncio from typing import TYPE_CHECKING from cloud.internal_api.models import DeviceSnapshotModel, HeartbeatResponse +from core.errors import DeviceRuntimeError from device.manager import DeviceManager from host_agent.client import HostAgentClient from host_agent.config import HostAgentConfig @@ -17,7 +18,7 @@ def build_device_snapshot(manager: DeviceManager) -> list[DeviceSnapshotModel]: DeviceSnapshotModel( device_id=device.id, driver_type=device.driver_type, - status=device.status, + status="idle" if device.status == "busy" else device.status, capability_tags=list(device.capability_tags), ) for device in sorted(manager.list_devices(), key=lambda item: item.id) @@ -47,6 +48,7 @@ class HeartbeatSynchronizer: ) async def run(self, stop: asyncio.Event) -> None: + self.connect_devices() while not stop.is_set(): await self.sync_once() try: @@ -56,3 +58,12 @@ class HeartbeatSynchronizer: ) except TimeoutError: continue + + def connect_devices(self) -> None: + for device in self.manager.list_devices(): + if device.status != "idle": + continue + try: + self.manager.connect(device.id) + except DeviceRuntimeError: + continue diff --git a/apps/device-host-agent/tests/test_heartbeat.py b/apps/device-host-agent/tests/test_heartbeat.py index bb197cc..4c7d006 100644 --- a/apps/device-host-agent/tests/test_heartbeat.py +++ b/apps/device-host-agent/tests/test_heartbeat.py @@ -9,6 +9,11 @@ from host_agent.config import HostAgentConfig from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot +class ConnectableDriver: + def connect(self) -> None: + return None + + def _config() -> HostAgentConfig: return HostAgentConfig( control_plane_url="https://control.example", @@ -38,7 +43,7 @@ def test_build_device_snapshot_copies_complete_non_secret_state() -> None: 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].status == "idle" assert snapshot[1].capability_tags == ["android", "physical"] assert "must-not-leave-host" not in repr(snapshot) @@ -47,7 +52,7 @@ def test_heartbeat_synchronizer_runs_at_configured_interval_until_stopped() -> N manager = DeviceManager() manager.register_device( "device-a", - lambda: object(), # type: ignore[arg-type,return-value] + lambda: ConnectableDriver(), # type: ignore[arg-type,return-value] ) calls: list[list[str]] = [] @@ -75,3 +80,4 @@ 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"