461 lines
15 KiB
Python
461 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from datetime import UTC, datetime
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from cloud.internal_api.models import AssignmentModel, DeviceSnapshotModel
|
|
from host_agent.client import (
|
|
HostAgentClient,
|
|
HostAgentEnrollmentClient,
|
|
HostAgentAPIError,
|
|
HostTaskSubmissionUnknownError,
|
|
StaleLeaseError,
|
|
)
|
|
from host_agent.config import HostAgentConfig
|
|
|
|
|
|
def _config(**overrides) -> HostAgentConfig:
|
|
values = {
|
|
"control_plane_url": "https://control.example",
|
|
"host_id": "host-a",
|
|
"token": "host-secret",
|
|
"retry_backoff_seconds": 0.01,
|
|
"max_retry_backoff_seconds": 0.02,
|
|
"max_retry_attempts": 3,
|
|
}
|
|
values.update(overrides)
|
|
return HostAgentConfig(**values)
|
|
|
|
|
|
def _assignment() -> AssignmentModel:
|
|
return 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",
|
|
)
|
|
|
|
|
|
def test_client_sends_authenticated_heartbeat_and_claim() -> None:
|
|
requests: list[httpx.Request] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
if request.url.path.endswith("/heartbeat"):
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"host_id": "host-a",
|
|
"accepted_devices": 1,
|
|
"received_at": "2026-07-12T00:00:00Z",
|
|
},
|
|
)
|
|
return httpx.Response(200, json={"assignment": None, "timed_out": True})
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
heartbeat = await client.heartbeat(
|
|
[
|
|
DeviceSnapshotModel(
|
|
device_id="device-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
)
|
|
]
|
|
)
|
|
assignment = await client.claim()
|
|
assert heartbeat.accepted_devices == 1
|
|
assert assignment is None
|
|
|
|
asyncio.run(scenario())
|
|
assert [request.url.path for request in requests] == [
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
"/internal/v1/hosts/host-a/assignments/claim",
|
|
]
|
|
assert all(
|
|
request.headers["authorization"] == "Bearer host-secret" for request in requests
|
|
)
|
|
|
|
|
|
def test_transport_failures_retry_with_bounded_backoff() -> None:
|
|
attempts = 0
|
|
sleeps: list[float] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
if attempts < 3:
|
|
raise httpx.ConnectError("temporary", request=request)
|
|
return httpx.Response(200, json={"assignment": None, "timed_out": True})
|
|
|
|
async def sleep(delay: float) -> None:
|
|
sleeps.append(delay)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client, sleep=sleep)
|
|
assert await client.claim() is None
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 3
|
|
assert sleeps == [0.01, 0.02]
|
|
|
|
|
|
def test_stale_lease_response_raises_typed_error_without_retry() -> None:
|
|
attempts = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
return httpx.Response(
|
|
409,
|
|
json={"code": "stale_lease", "detail": "lease expired"},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
with pytest.raises(StaleLeaseError) as error:
|
|
await client.renew(_assignment())
|
|
assert "host-secret" not in str(error.value)
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|
|
|
|
|
|
def test_renew_deserializes_cancel_requested_flag() -> None:
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"status": "renewed",
|
|
"lease_expires_at": "2026-07-12T00:05:00Z",
|
|
"cancel_requested": True,
|
|
},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
response = await client.renew(_assignment())
|
|
assert response.status == "renewed"
|
|
assert response.cancel_requested is True
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_result_report_retries_identical_payload_after_response_loss() -> None:
|
|
payloads: list[dict[str, object]] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
payloads.append(json.loads(request.content))
|
|
if len(payloads) == 1:
|
|
raise httpx.ReadError("response lost", request=request)
|
|
return httpx.Response(200, json={"status": "already_recorded"})
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(
|
|
_config(),
|
|
http_client=http_client,
|
|
sleep=lambda delay: asyncio.sleep(0),
|
|
)
|
|
response = await client.report_result(
|
|
_assignment(),
|
|
status="failed",
|
|
failure_reason="planner unavailable",
|
|
result={"runtime_status": "failed"},
|
|
)
|
|
|
|
assert response.status == "already_recorded"
|
|
|
|
asyncio.run(scenario())
|
|
assert len(payloads) == 2
|
|
assert payloads[0] == payloads[1]
|
|
assert payloads[0]["failure_reason"] == "planner unavailable"
|
|
|
|
|
|
def test_submit_self_task_posts_once_and_returns_task_id() -> None:
|
|
requests: list[httpx.Request] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
requests.append(request)
|
|
return httpx.Response(201, json={"task_id": "task-cloud-1"})
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
response = await client.submit_self_task(
|
|
goal="open settings",
|
|
device_id=None,
|
|
)
|
|
|
|
assert response.task_id == "task-cloud-1"
|
|
|
|
asyncio.run(scenario())
|
|
assert len(requests) == 1
|
|
assert requests[0].url.path == "/internal/v1/hosts/host-a/tasks"
|
|
assert requests[0].headers["authorization"] == "Bearer host-secret"
|
|
payload = json.loads(requests[0].content)
|
|
assert payload == {
|
|
"host_id": "host-a",
|
|
"goal": "open settings",
|
|
"device_id": None,
|
|
}
|
|
|
|
|
|
def test_submit_self_task_does_not_retry_transport_failure() -> None:
|
|
attempts = 0
|
|
sleeps: list[float] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
raise httpx.ConnectError("network down", request=request)
|
|
|
|
async def sleep(delay: float) -> None:
|
|
sleeps.append(delay)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client, sleep=sleep)
|
|
with pytest.raises(HostTaskSubmissionUnknownError):
|
|
await client.submit_self_task(goal="open settings")
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|
|
assert sleeps == []
|
|
|
|
|
|
def test_submit_self_task_treats_5xx_as_unknown_outcome_without_retry() -> None:
|
|
attempts = 0
|
|
sleeps: list[float] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
return httpx.Response(502, json={"detail": "bad gateway"})
|
|
|
|
async def sleep(delay: float) -> None:
|
|
sleeps.append(delay)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client, sleep=sleep)
|
|
with pytest.raises(HostTaskSubmissionUnknownError):
|
|
await client.submit_self_task(goal="open settings")
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|
|
assert sleeps == []
|
|
|
|
|
|
def test_submit_self_task_raises_definitive_error_on_4xx_rejection() -> None:
|
|
attempts = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
return httpx.Response(
|
|
403,
|
|
json={"detail": "Host self-submission is disabled"},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
with pytest.raises(HostAgentAPIError) as error:
|
|
await client.submit_self_task(goal="open settings")
|
|
assert "host-secret" not in str(error.value)
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|
|
|
|
|
|
def test_submit_self_task_treats_malformed_success_as_unknown() -> None:
|
|
attempts = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
return httpx.Response(201, json={"unexpected": "shape"})
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
with pytest.raises(HostTaskSubmissionUnknownError):
|
|
await client.submit_self_task(goal="open settings")
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|
|
|
|
|
|
def test_submit_self_task_does_not_duplicate_when_response_is_lost() -> None:
|
|
attempts = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
raise httpx.ReadError("response lost", request=request)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(
|
|
_config(),
|
|
http_client=http_client,
|
|
sleep=lambda delay: asyncio.sleep(0),
|
|
)
|
|
with pytest.raises(HostTaskSubmissionUnknownError):
|
|
await client.submit_self_task(goal="open settings")
|
|
|
|
asyncio.run(scenario())
|
|
assert attempts == 1
|
|
|
|
|
|
def test_heartbeat_includes_mcp_busy_device_ids_in_payload() -> None:
|
|
"""When mcp_busy_device_ids is passed, the client sends it in the request."""
|
|
captured: list[dict[str, object]] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
captured.append(json.loads(request.content))
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"host_id": "host-a",
|
|
"accepted_devices": 0,
|
|
"received_at": "2026-07-12T00:00:00Z",
|
|
},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
await client.heartbeat([], mcp_busy_device_ids=["phone-1"])
|
|
|
|
asyncio.run(scenario())
|
|
assert captured[0]["mcp_busy_device_ids"] == ["phone-1"]
|
|
|
|
|
|
def test_heartbeat_omits_mcp_busy_device_ids_when_empty() -> None:
|
|
"""Empty list is omitted from the payload (backward compatible)."""
|
|
captured: list[dict[str, object]] = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
captured.append(json.loads(request.content))
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"host_id": "host-a",
|
|
"accepted_devices": 0,
|
|
"received_at": "2026-07-12T00:00:00Z",
|
|
},
|
|
)
|
|
|
|
async def scenario() -> None:
|
|
async with httpx.AsyncClient(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentClient(_config(), http_client=http_client)
|
|
await client.heartbeat([], mcp_busy_device_ids=[])
|
|
|
|
asyncio.run(scenario())
|
|
assert "mcp_busy_device_ids" not in captured[0]
|
|
|
|
|
|
def test_bootstrap_client_directly_enrolls_and_enrolls_device() -> None:
|
|
requests: list[httpx.Request] = []
|
|
host_attempts = 0
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
nonlocal host_attempts
|
|
requests.append(request)
|
|
if request.url.path == "/internal/v1/enrollments":
|
|
host_attempts += 1
|
|
if host_attempts == 1:
|
|
raise httpx.ReadError("response lost", request=request)
|
|
return httpx.Response(201, json={"host_id": "host-cloud-a"})
|
|
return httpx.Response(201, json={"device_id": "device-cloud-a"})
|
|
|
|
config = _config(
|
|
host_id="",
|
|
token="",
|
|
enrollment_managed=True,
|
|
)
|
|
with httpx.Client(
|
|
transport=httpx.MockTransport(handler),
|
|
base_url="https://control.example",
|
|
) as http_client:
|
|
client = HostAgentEnrollmentClient(
|
|
config,
|
|
http_client=http_client,
|
|
sleep=lambda _delay: None,
|
|
)
|
|
host = client.enroll_host(
|
|
agent_instance_id="agent-instance-a",
|
|
host_token="host-token-" + ("x" * 40),
|
|
display_name="Edge Mac",
|
|
)
|
|
client.config = _config(
|
|
host_id=host.host_id,
|
|
token="host-token-" + ("x" * 40),
|
|
enrollment_managed=True,
|
|
)
|
|
device = client.enroll_device(
|
|
local_device_id="local-device-a",
|
|
driver_type="wda",
|
|
name="iPhone",
|
|
capability_tags=["ios"],
|
|
)
|
|
|
|
assert host.host_id == "host-cloud-a"
|
|
assert device.device_id == "device-cloud-a"
|
|
assert len(requests) == 3
|
|
assert requests[0].content == requests[1].content
|
|
assert "authorization" not in requests[0].headers
|
|
assert requests[2].headers["authorization"] == ("Bearer host-token-" + ("x" * 40))
|