feat(host-protocol): define internal API models

2026-07-12 18:07:42 +08:00
parent a05465458c
commit b4cbd83dcd
4 changed files with 177 additions and 1 deletions
@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+import pytest
+from pydantic import ValidationError
+
+from cloud.internal_api.models import (
+    AssignmentModel,
+    ClaimRequest,
+    ClaimResponse,
+    DeviceSnapshotModel,
+    HeartbeatRequest,
+    LeaseRenewalRequest,
+    TerminalResultRequest,
+)
+
+
+def test_heartbeat_models_complete_device_snapshot() -> None:
+    heartbeat = HeartbeatRequest(
+        host_id="host-a",
+        address="10.0.0.1:9000",
+        devices=[
+            DeviceSnapshotModel(
+                device_id="device-a",
+                driver_type="wda",
+                status="idle",
+                capability_tags=["ios", "physical"],
+            )
+        ],
+    )
+
+    assert heartbeat.host_id == "host-a"
+    assert heartbeat.devices[0].capability_tags == ["ios", "physical"]
+
+
+@pytest.mark.parametrize("status", ["unreachable", "unknown", ""])
+def test_device_snapshot_rejects_invalid_host_status(status: str) -> None:
+    with pytest.raises(ValidationError):
+        DeviceSnapshotModel(
+            device_id="device-a",
+            driver_type="wda",
+            status=status,
+        )
+
+
+def test_claim_response_supports_assignment_and_normal_timeout() -> None:
+    assignment = AssignmentModel(
+        task_id="task-a",
+        attempt=1,
+        lease_id="lease-a",
+        lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC),
+        host_id="host-a",
+        device_id="device-a",
+        goal="open settings",
+    )
+
+    assert ClaimResponse(assignment=assignment).assignment == assignment
+    assert ClaimResponse(timed_out=True).model_dump() == {
+        "assignment": None,
+        "timed_out": True,
+    }
+
+
+@pytest.mark.parametrize("timeout_seconds", [-1, 61])
+def test_claim_timeout_is_bounded(timeout_seconds: float) -> None:
+    with pytest.raises(ValidationError):
+        ClaimRequest(host_id="host-a", timeout_seconds=timeout_seconds)
+
+
+def test_renewal_and_result_require_attempt_and_lease_identity() -> None:
+    renewal = LeaseRenewalRequest(
+        host_id="host-a",
+        task_id="task-a",
+        attempt=2,
+        lease_id="lease-a",
+    )
+    result = TerminalResultRequest(
+        host_id="host-a",
+        task_id="task-a",
+        attempt=2,
+        lease_id="lease-a",
+        status="failed",
+        failure_reason="device offline",
+        result={"step": 3},
+    )
+
+    assert renewal.attempt == result.attempt == 2
+    assert result.failure_reason == "device offline"
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        {"host_id": "", "task_id": "task-a", "attempt": 1, "lease_id": "lease"},
+        {"host_id": "host-a", "task_id": "", "attempt": 1, "lease_id": "lease"},
+        {"host_id": "host-a", "task_id": "task-a", "attempt": 0, "lease_id": "lease"},
+        {"host_id": "host-a", "task_id": "task-a", "attempt": 1, "lease_id": ""},
+    ],
+)
+def test_renewal_rejects_incomplete_identity(payload: dict[str, object]) -> None:
+    with pytest.raises(ValidationError):
+        LeaseRenewalRequest.model_validate(payload)