agentic-mobile-control/apps/device-host-agent/tests/test_client.py

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,
    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_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_bootstrap_client_retries_identical_enrollment_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_token="one-time-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_token="one-time-token",
            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 requests[0].headers["authorization"] == "Bearer one-time-token"
    assert requests[2].headers["authorization"] == ("Bearer host-token-" + ("x" * 40))