agentic-mobile-control/apps/device-host-agent/tests/test_devices.py

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() == []