104 lines
3.0 KiB
Python
104 lines
3.0 KiB
Python
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)
|