from __future__ import annotations from cloud.internal_api.models import DeviceEnrollmentResponse from device.manager import DeviceManager from host_agent.config import HostAgentConfig from host_agent.devices import register_local_device, unregister_local_device from storage.device_config import DeviceConfigStore def _config(*, enrollment_managed: bool) -> HostAgentConfig: return HostAgentConfig( control_plane_url="https://control.example", host_id="host-a", token="secret", enrollment_managed=enrollment_managed, ) class RecordingEnrollmentClient: def __init__(self, device_id: str) -> None: self.device_id = device_id self.calls: list[dict[str, object]] = [] def enroll_device(self, **payload): self.calls.append(payload) return DeviceEnrollmentResponse(device_id=self.device_id) def test_register_local_device_enrollment_managed_registers_under_cloud_id( tmp_path, ) -> None: store = DeviceConfigStore(tmp_path / "devices.sqlite3") manager = DeviceManager() enrollment_client = RecordingEnrollmentClient("device-cloud-a") register_local_device( store, manager, device_id="local-device-a", driver_type="wda", connection_info={"server_url": "http://127.0.0.1:4723"}, name="Lab iPhone", config=_config(enrollment_managed=True), enrollment_client=enrollment_client, # type: ignore[arg-type] ) assert enrollment_client.calls == [ { "local_device_id": "local-device-a", "driver_type": "wda", "name": "Lab iPhone", "capability_tags": [], } ] assert [device.id for device in manager.list_devices()] == ["device-cloud-a"] assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a" def test_register_local_device_not_enrollment_managed_registers_under_local_id( tmp_path, ) -> None: store = DeviceConfigStore(tmp_path / "devices.sqlite3") manager = DeviceManager() enrollment_client = RecordingEnrollmentClient("device-cloud-a") register_local_device( store, manager, device_id="local-device-a", driver_type="wda", connection_info={"server_url": "http://127.0.0.1:4723"}, name="Lab iPhone", config=_config(enrollment_managed=False), enrollment_client=enrollment_client, # type: ignore[arg-type] ) assert enrollment_client.calls == [] assert [device.id for device in manager.list_devices()] == ["local-device-a"] assert store.get("local-device-a")["cloud_device_id"] is None def test_register_local_device_reregister_under_new_cloud_id_replaces_prior_entry( tmp_path, ) -> None: store = DeviceConfigStore(tmp_path / "devices.sqlite3") manager = DeviceManager() config = _config(enrollment_managed=True) register_local_device( store, manager, device_id="local-device-a", driver_type="wda", connection_info={"server_url": "http://127.0.0.1:4723"}, name="Lab iPhone", config=config, enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type] "device-cloud-a" ), ) assert [device.id for device in manager.list_devices()] == ["device-cloud-a"] register_local_device( store, manager, device_id="local-device-a", driver_type="wda", connection_info={"server_url": "http://127.0.0.1:5000"}, name="Lab iPhone (moved)", config=config, enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type] "device-cloud-b" ), ) devices = manager.list_devices() assert [device.id for device in devices] == ["device-cloud-b"] assert devices[0].name == "Lab iPhone (moved)" assert devices[0].connection_info == {"server_url": "http://127.0.0.1:5000"} assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-b" def test_unregister_local_device_removes_device_with_cloud_id(tmp_path) -> None: store = DeviceConfigStore(tmp_path / "devices.sqlite3") manager = DeviceManager() register_local_device( store, manager, device_id="local-device-a", driver_type="wda", connection_info={}, name=None, config=_config(enrollment_managed=True), enrollment_client=RecordingEnrollmentClient( # type: ignore[arg-type] "device-cloud-a" ), ) unregister_local_device(store, manager, device_id="local-device-a") assert manager.list_devices() == [] assert store.get("local-device-a") is None def test_unregister_local_device_removes_device_without_cloud_id(tmp_path) -> None: store = DeviceConfigStore(tmp_path / "devices.sqlite3") manager = DeviceManager() register_local_device( store, manager, device_id="local-device-a", driver_type="wda", connection_info={}, name=None, config=_config(enrollment_managed=False), enrollment_client=None, ) unregister_local_device(store, manager, device_id="local-device-a") assert manager.list_devices() == [] assert store.get("local-device-a") is None def test_unregister_local_device_is_a_no_op_for_unknown_device(tmp_path) -> None: store = DeviceConfigStore(tmp_path / "devices.sqlite3") manager = DeviceManager() unregister_local_device(store, manager, device_id="unknown-device") assert manager.list_devices() == []