feat(cloud): add edge host enrollment

apps/device-host-agent/tests/test_app.py

@@ -4,10 +4,15 @@ import asyncio
 from contextlib import suppress
 from datetime import UTC, datetime, timedelta
 
-from cloud.internal_api.models import AssignmentModel
+from cloud.internal_api.models import (
+    AssignmentModel,
+    DeviceEnrollmentResponse,
+    HostEnrollmentResponse,
+)
 from device.manager import DeviceManager
 from host_agent.app import HostAgentApplication, create_application
 from host_agent.config import HostAgentConfig
+from host_agent.identity import HostIdentityStore
 from storage.device_config import DeviceConfigStore
 
 
@@ -65,6 +70,106 @@ def test_create_application_loads_persisted_device_configuration(
     asyncio.run(application.client.aclose())
 
 
+def test_create_application_enrolls_host_and_devices_before_managed_startup(
+    tmp_path,
+    monkeypatch,
+) -> None:
+    monkeypatch.chdir(tmp_path)
+    store = DeviceConfigStore(tmp_path / "devices.sqlite3")
+    store.add(
+        device_id="local-device-a",
+        name="Lab iPhone",
+        driver_type="wda",
+        connection_info={"server_url": "http://127.0.0.1:4723"},
+    )
+    events: list[str] = []
+
+    class EnrollmentClient:
+        def __init__(self) -> None:
+            self.config = HostAgentConfig(
+                control_plane_url="https://control.example",
+                enrollment_token="one-time-token",
+                enrollment_managed=True,
+            )
+
+        def enroll_host(self, **payload):
+            events.append(f"host:{payload['agent_instance_id']}")
+            return HostEnrollmentResponse(host_id="host-cloud-a")
+
+        def enroll_device(self, **payload):
+            events.append(f"device:{payload['local_device_id']}")
+            return DeviceEnrollmentResponse(device_id="device-cloud-a")
+
+        def close(self):
+            raise AssertionError("injected client must not be closed")
+
+    identity_store = HostIdentityStore(tmp_path / "host_identity.json")
+    enrollment_client = EnrollmentClient()
+    application = create_application(
+        config=enrollment_client.config,
+        device_config_store=store,
+        identity_store=identity_store,
+        enrollment_client=enrollment_client,  # type: ignore[arg-type]
+    )
+
+    assert events[0].startswith("host:agent-")
+    assert events[1] == "device:local-device-a"
+    assert application.client.config.host_id == "host-cloud-a"
+    assert application.client.config.enrollment_managed is True
+    assert [device.id for device in application.heartbeat.manager.list_devices()] == [
+        "device-cloud-a"
+    ]
+    assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a"
+    assert identity_store.load().host_id == "host-cloud-a"
+    asyncio.run(application.client.aclose())
+
+
+def test_managed_restart_reuses_identity_and_recovers_device_mapping(
+    tmp_path,
+    monkeypatch,
+) -> None:
+    monkeypatch.chdir(tmp_path)
+    store = DeviceConfigStore(tmp_path / "devices.sqlite3")
+    store.add(
+        device_id="local-device-a",
+        driver_type="wda",
+        connection_info={},
+    )
+    identity_store = HostIdentityStore(tmp_path / "host_identity.json")
+    identity_store.complete(identity_store.load_or_create(), "host-cloud-a")
+    events: list[str] = []
+
+    class EnrollmentClient:
+        config = HostAgentConfig(
+            control_plane_url="https://control.example",
+            identity_path=tmp_path / "host_identity.json",
+            enrollment_managed=True,
+        )
+
+        def enroll_host(self, **payload):
+            raise AssertionError("completed identity must skip Host enrollment")
+
+        def enroll_device(self, **payload):
+            events.append(payload["local_device_id"])
+            return DeviceEnrollmentResponse(device_id="device-cloud-a")
+
+        def close(self):
+            return None
+
+    enrollment_client = EnrollmentClient()
+    application = create_application(
+        config=enrollment_client.config,
+        device_config_store=store,
+        identity_store=identity_store,
+        enrollment_client=enrollment_client,  # type: ignore[arg-type]
+    )
+
+    assert events == ["local-device-a"]
+    assert application.client.config.host_id == "host-cloud-a"
+    assert store.get("local-device-a")["cloud_device_id"] == "device-cloud-a"
+    asyncio.run(application.client.aclose())
+
+
 def test_shutdown_cancels_long_poll_and_sends_final_heartbeat() -> None:
     async def scenario() -> None:
         claim_started = asyncio.Event()