agentic-mobile-control/tests/test_host_agent_progress_reporting.py

"""Tests for Host Agent progress reporting via lease renewal (task 2.4)."""

from __future__ import annotations

import asyncio
from datetime import UTC, datetime, timedelta
from threading import Event
from typing import Any

import httpx

from cloud.internal_api.models import (
    AssignmentModel,
    LeaseRenewalResponse,
)
from host_agent.config import HostAgentConfig
from host_agent.lease import ActiveAssignmentRunner
from host_agent.progress import TaskProgressHolder, TaskProgressSnapshot


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #


def _assignment(
    *,
    lease_expires_at: datetime | None = None,
) -> AssignmentModel:
    return AssignmentModel(
        task_id="task-progress",
        attempt=1,
        lease_id="lease-1",
        lease_expires_at=lease_expires_at or datetime.now(UTC) + timedelta(minutes=5),
        host_id="host-a",
        device_id="device-a",
        goal="do something",
    )


class _FakeExecutor:
    """Stub executor that optionally fires a progress callback then blocks."""

    def __init__(
        self,
        *,
        fire_progress_before_block: bool = False,
        progress_payload: tuple[int, str, str] = (1, "running", "step 1 summary"),
        block_event: Event | None = None,
    ) -> None:
        self._fire = fire_progress_before_block
        self._payload = progress_payload
        self._block_event = block_event or Event()
        self._progress = TaskProgressHolder()

    def execute(
        self,
        assignment: AssignmentModel,
        *,
        should_stop: Any | None = None,
        stop_reason: Any | None = None,
    ) -> Any:
        from host_agent.assignment import AssignmentExecutionResult

        if self._fire:
            self._progress.update(*self._payload)

        # Block until test signals completion.
        self._block_event.wait(timeout=5)
        return AssignmentExecutionResult(status="done")

    def latest_progress(self) -> TaskProgressSnapshot | None:
        return self._progress.snapshot()


class _RecordingClient:
    """Fake HostAgentClient that records each renew() call's progress argument."""

    def __init__(self) -> None:
        self.renew_calls: list[TaskProgressSnapshot | None] = []
        self._call_count = 0

    async def renew(
        self,
        assignment: AssignmentModel,
        *,
        progress: TaskProgressSnapshot | None = None,
    ) -> LeaseRenewalResponse:
        self.renew_calls.append(progress)
        self._call_count += 1
        return LeaseRenewalResponse(
            status="renewed",
            lease_expires_at=datetime.now(UTC) + timedelta(minutes=5),
        )


def _runner(
    client: _RecordingClient, executor: _FakeExecutor
) -> ActiveAssignmentRunner:
    return ActiveAssignmentRunner(client, executor, now=lambda: datetime.now(UTC))


# --------------------------------------------------------------------------- #
# Renewal includes progress after a step completes
# --------------------------------------------------------------------------- #


def test_renewal_includes_progress_after_step(tmp_path) -> None:
    block_event = Event()
    executor = _FakeExecutor(
        fire_progress_before_block=True,
        progress_payload=(1, "running", "step 1 summary"),
        block_event=block_event,
    )
    client = _RecordingClient()
    runner = _runner(client, executor)

    # Set a very short lease so renewal fires quickly.
    assignment = _assignment(
        lease_expires_at=datetime.now(UTC) + timedelta(milliseconds=50),
    )

    async def _drive() -> None:
        # Start the runner; execution thread will fire progress then block.
        task = asyncio.create_task(runner.run(assignment))
        # Wait for at least one renewal to happen.
        for _ in range(50):
            await asyncio.sleep(0.05)
            if client.renew_calls:
                break
        # Signal the executor to complete.
        block_event.set()
        await task

    asyncio.run(_drive())

    assert len(client.renew_calls) > 0
    # At least one renewal should have non-None progress.
    progress_renewals = [p for p in client.renew_calls if p is not None]
    assert len(progress_renewals) > 0
    snap = progress_renewals[0]
    assert snap.step_index == 1
    assert snap.step_status == "running"
    assert "step 1 summary" in snap.summary


# --------------------------------------------------------------------------- #
# Renewal omits progress before any step completes
# --------------------------------------------------------------------------- #


def test_renewal_omits_progress_before_any_step(tmp_path) -> None:
    block_event = Event()
    executor = _FakeExecutor(
        fire_progress_before_block=False,
        block_event=block_event,
    )
    client = _RecordingClient()
    runner = _runner(client, executor)

    assignment = _assignment(
        lease_expires_at=datetime.now(UTC) + timedelta(milliseconds=50),
    )

    async def _drive() -> None:
        task = asyncio.create_task(runner.run(assignment))
        for _ in range(50):
            await asyncio.sleep(0.05)
            if client.renew_calls:
                break
        block_event.set()
        await task

    asyncio.run(_drive())

    assert len(client.renew_calls) > 0
    # No progress should have been reported.
    assert all(p is None for p in client.renew_calls)


# --------------------------------------------------------------------------- #
# Oversized summary truncated client-side before sending
# --------------------------------------------------------------------------- #


def test_oversized_summary_truncated_client_side() -> None:
    """HostAgentClient.renew truncates summary to <=500 chars before serializing."""
    from host_agent.client import HostAgentClient

    captured_payload: dict[str, Any] = {}

    def _handler(request: httpx.Request) -> httpx.Response:
        import json

        body = json.loads(request.content.decode("utf-8"))
        captured_payload.update(body)
        return httpx.Response(
            200,
            json={
                "status": "renewed",
                "lease_expires_at": (
                    datetime.now(UTC) + timedelta(minutes=5)
                ).isoformat(),
            },
        )

    transport = httpx.MockTransport(_handler)
    config = HostAgentConfig(
        control_plane_url="https://control.example",
        host_id="host-a",
        token="secret",
        retry_backoff_seconds=0.01,
        max_retry_attempts=1,
    )
    http_client = httpx.AsyncClient(
        base_url=config.control_plane_url, transport=transport
    )
    client = HostAgentClient(config, http_client=http_client)

    long_summary = "x" * 2000
    snapshot = TaskProgressSnapshot(
        step_index=1,
        step_status="running",
        summary=long_summary,
        updated_at=datetime.now(UTC),
    )

    asyncio.run(client.renew(_assignment(), progress=snapshot))

    assert "progress" in captured_payload
    assert captured_payload["progress"] is not None
    serialized_summary = captured_payload["progress"]["summary"]
    assert len(serialized_summary) <= 500
    assert serialized_summary == "x" * 500

    asyncio.run(http_client.aclose())