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

from __future__ import annotations

import asyncio
from datetime import UTC, datetime

from cloud.internal_api.models import AssignmentModel, TerminalResultResponse
from host_agent.assignment import AssignmentExecutionResult
from host_agent.processor import AssignmentProcessor
from host_agent.status import AgentStatusTracker


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_processor_reports_success_with_execution_metadata() -> None:
    async def scenario() -> None:
        reports: list[dict[str, object]] = []

        class SuccessfulExecutor:
            async def run(self, assignment):
                return AssignmentExecutionResult(
                    status="done",
                    failure_reason="ignored success detail",
                    metadata={"runtime_task_id": "runtime-a"},
                )

        class RecordingClient:
            async def report_result(self, assignment, **kwargs):
                reports.append(kwargs)
                return TerminalResultResponse(status="recorded")

        result = await AssignmentProcessor(
            RecordingClient(),  # type: ignore[arg-type]
            SuccessfulExecutor(),
        ).process(_assignment())

        assert result.report_status == "recorded"
        assert reports == [
            {
                "status": "done",
                "failure_reason": None,
                "result": {"runtime_task_id": "runtime-a"},
            }
        ]

    asyncio.run(scenario())


def test_processor_preserves_runtime_failure_reason() -> None:
    async def scenario() -> None:
        reports: list[dict[str, object]] = []

        class FailedExecutor:
            async def run(self, assignment):
                return AssignmentExecutionResult(
                    status="failed",
                    failure_reason="planner unavailable",
                    metadata={"runtime_status": "failed"},
                )

        class RecordingClient:
            async def report_result(self, assignment, **kwargs):
                reports.append(kwargs)
                return TerminalResultResponse(status="already_recorded")

        result = await AssignmentProcessor(
            RecordingClient(),  # type: ignore[arg-type]
            FailedExecutor(),
        ).process(_assignment())

        assert result.report_status == "already_recorded"
        assert result.execution.failure_reason == "planner unavailable"
        assert reports[0] == {
            "status": "failed",
            "failure_reason": "planner unavailable",
            "result": {"runtime_status": "failed"},
        }

    asyncio.run(scenario())


def test_processor_reports_cancelled_status_with_reason() -> None:
    async def scenario() -> None:
        reports: list[dict[str, object]] = []

        class CancelledExecutor:
            async def run(self, assignment):
                return AssignmentExecutionResult(
                    status="cancelled",
                    failure_reason="cancellation requested by control plane",
                    metadata={"runtime_status": "cancelled"},
                )

        class RecordingClient:
            async def report_result(self, assignment, **kwargs):
                reports.append(kwargs)
                return TerminalResultResponse(status="recorded")

        result = await AssignmentProcessor(
            RecordingClient(),  # type: ignore[arg-type]
            CancelledExecutor(),
        ).process(_assignment())

        assert result.report_status == "recorded"
        assert reports[0] == {
            "status": "cancelled",
            "failure_reason": "cancellation requested by control plane",
            "result": {"runtime_status": "cancelled"},
        }

    asyncio.run(scenario())


def test_status_tracker_sees_started_then_finished_even_on_raise() -> None:
    async def scenario() -> None:
        tracker = AgentStatusTracker()
        snapshots: list[dict[str, object]] = []

        class RaisingExecutor:
            async def run(self, assignment):
                snapshots.append(tracker.snapshot())
                raise RuntimeError("executor exploded")

        class RecordingClient:
            async def report_result(self, assignment, **kwargs):
                return TerminalResultResponse(status="recorded")

        processor = AssignmentProcessor(
            RecordingClient(),  # type: ignore[arg-type]
            RaisingExecutor(),
            status_tracker=tracker,
        )

        try:
            await processor.process(_assignment())
        except RuntimeError:
            pass

        assert snapshots[0]["current_assignment"] is not None
        assert snapshots[0]["current_assignment"]["task_id"] == "task-a"
        assert tracker.snapshot()["current_assignment"] is None

    asyncio.run(scenario())


def test_on_result_receives_assignment_and_result_and_swallows_exceptions() -> None:
    async def scenario() -> None:
        received: list[tuple[object, object]] = []

        class SuccessfulExecutor:
            async def run(self, assignment):
                return AssignmentExecutionResult(
                    status="done",
                    failure_reason=None,
                    metadata={},
                )

        class RecordingClient:
            async def report_result(self, assignment, **kwargs):
                return TerminalResultResponse(status="recorded")

        def on_result(assignment, result) -> None:
            received.append((assignment, result))
            raise RuntimeError("history recording exploded")

        assignment = _assignment()
        processor = AssignmentProcessor(
            RecordingClient(),  # type: ignore[arg-type]
            SuccessfulExecutor(),
            on_result=on_result,
        )

        result = await processor.process(assignment)

        assert received == [(assignment, result)]

    asyncio.run(scenario())