139 lines
4.2 KiB
Python
139 lines
4.2 KiB
Python
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
|