feat(host-protocol): accept heartbeat snapshots

2026-07-12 18:10:42 +08:00
parent b4cbd83dcd
commit efb754fbe7
3 changed files with 246 additions and 1 deletions
@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from cloud.auth import BearerCredential, ConfiguredBearerAuthProvider
+from cloud.config import CloudConfig
+from cloud.internal_api.api import create_internal_router
+from cloud.pool import DevicePool, PooledDevice
+from cloud.store import CloudStore
+
+
+def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
+    pool = DevicePool(
+        CloudStore(tmp_path / "internal.sqlite3"),
+        CloudConfig(stale_after_seconds=60),
+    )
+    auth_provider = ConfiguredBearerAuthProvider(
+        [
+            BearerCredential(
+                principal_id="agent-a",
+                token="token-a",
+                host_id="host-a",
+            ),
+            BearerCredential(
+                principal_id="agent-b",
+                token="token-b",
+                host_id="host-b",
+            ),
+        ]
+    )
+    app = FastAPI()
+    app.include_router(create_internal_router(pool=pool, auth_provider=auth_provider))
+    return TestClient(app), pool
+
+
+def _heartbeat_payload(host_id: str, *device_ids: str) -> dict[str, object]:
+    return {
+        "host_id": host_id,
+        "address": f"{host_id}.internal",
+        "devices": [
+            {
+                "device_id": device_id,
+                "driver_type": "wda",
+                "status": "idle",
+                "capability_tags": ["ios"],
+            }
+            for device_id in device_ids
+        ],
+    }
+
+
+def test_authenticated_heartbeat_replaces_complete_snapshot(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+
+    first = client.put(
+        "/internal/v1/hosts/host-a/heartbeat",
+        headers={"Authorization": "Bearer token-a"},
+        json=_heartbeat_payload("host-a", "device-1", "device-2"),
+    )
+    second = client.put(
+        "/internal/v1/hosts/host-a/heartbeat",
+        headers={"Authorization": "Bearer token-a"},
+        json=_heartbeat_payload("host-a", "device-2", "device-3"),
+    )
+
+    assert first.status_code == 200
+    assert second.status_code == 200
+    assert second.json()["accepted_devices"] == 2
+    assert {device.device_id for device in pool.store.list_devices()} == {
+        "device-2",
+        "device-3",
+    }
+
+
+def test_invalid_duplicate_snapshot_preserves_previous_devices(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    client.put(
+        "/internal/v1/hosts/host-a/heartbeat",
+        headers={"Authorization": "Bearer token-a"},
+        json=_heartbeat_payload("host-a", "existing-device"),
+    )
+
+    response = client.put(
+        "/internal/v1/hosts/host-a/heartbeat",
+        headers={"Authorization": "Bearer token-a"},
+        json=_heartbeat_payload("host-a", "duplicate", "duplicate"),
+    )
+
+    assert response.status_code == 422
+    assert [device.device_id for device in pool.store.list_devices()] == [
+        "existing-device"
+    ]
+
+
+def test_live_device_owner_conflict_is_rejected_without_partial_sync(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+    now = datetime.now(UTC)
+    pool.store.upsert_host("host-a", address=None, last_seen_at=now)
+    pool.store.replace_host_devices(
+        "host-a",
+        [
+            PooledDevice(
+                device_id="shared-device",
+                host_id="host-a",
+                driver_type="wda",
+                status="idle",
+                synced_at=now,
+            )
+        ],
+    )
+
+    response = client.put(
+        "/internal/v1/hosts/host-b/heartbeat",
+        headers={"Authorization": "Bearer token-b"},
+        json=_heartbeat_payload("host-b", "shared-device"),
+    )
+
+    assert response.status_code == 409
+    assert pool.store.get_host("host-b") is None
+    devices = pool.store.list_devices()
+    assert len(devices) == 1
+    assert devices[0].host_id == "host-a"
+
+
+def test_host_token_cannot_submit_heartbeat_for_another_host(tmp_path) -> None:
+    client, pool = _build_client(tmp_path)
+
+    response = client.put(
+        "/internal/v1/hosts/host-b/heartbeat",
+        headers={"Authorization": "Bearer token-a"},
+        json=_heartbeat_payload("host-b", "device-b"),
+    )
+
+    assert response.status_code == 403
+    assert pool.store.get_host("host-b") is None