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

from __future__ import annotations

from datetime import UTC, datetime
from types import SimpleNamespace

from cloud.internal_api.models import AssignmentModel
from core.models import Task
from host_agent.assignment import AssignmentExecutor
from host_agent.execution import ExecutionFactories


def _assignment(**overrides) -> AssignmentModel:
    values = {
        "task_id": "cloud-task",
        "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",
    }
    values.update(overrides)
    return AssignmentModel(**values)


def test_goal_assignment_executes_through_task_runner() -> None:
    received: list[Task] = []
    created: list[tuple[str, str | None, int | None]] = []

    class FakeTaskRunner:
        def run(self, task: Task) -> Task:
            received.append(task)
            task.status = "completed"
            return task

    class FakeMetadataStore:
        def create_task(
            self,
            task: Task,
            *,
            source_task_id: str | None = None,
            source_attempt: int | None = None,
        ) -> None:
            created.append((task.id, source_task_id, source_attempt))

    factories = ExecutionFactories(
        task_runner_factory=lambda: FakeTaskRunner(),  # type: ignore[arg-type,return-value]
        workflow_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_store=object(),  # type: ignore[arg-type]
        metadata_store=FakeMetadataStore(),  # type: ignore[arg-type]
    )

    result = AssignmentExecutor(factories).execute(_assignment())

    assert result.status == "done"
    assert received[0].goal == "open settings"
    assert received[0].device_id == "device-a"
    assert result.metadata["runtime_task_id"] == received[0].id
    assert created == [(received[0].id, "cloud-task", 1)]


def test_goal_assignment_preserves_runtime_failure_reason() -> None:
    class FakeTaskRunner:
        def run(self, task: Task) -> Task:
            task.status = "failed"
            task.failure_reason = "planner unavailable"
            return task

    factories = ExecutionFactories(
        task_runner_factory=lambda: FakeTaskRunner(),  # type: ignore[arg-type,return-value]
        workflow_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_store=object(),  # type: ignore[arg-type]
    )

    result = AssignmentExecutor(factories).execute(_assignment())

    assert result.status == "failed"
    assert result.failure_reason == "planner unavailable"


def test_goal_assignment_maps_cancellation_stop_to_cancelled_status() -> None:
    # First should_stop() call is Executor.execute()'s pre-flight check (must pass
    # through so the runner is actually invoked); the runner's own loop then stops.
    calls = {"count": 0}

    def should_stop() -> bool:
        calls["count"] += 1
        return calls["count"] > 1

    class FakeTaskRunner:
        def run(self, task: Task, *, should_stop=None, stop_reason=None) -> Task:
            assert should_stop is not None and should_stop()
            assert stop_reason is not None
            task.status = "cancelled"
            task.failure_reason = stop_reason()
            return task

    factories = ExecutionFactories(
        task_runner_factory=lambda: FakeTaskRunner(),  # type: ignore[arg-type,return-value]
        workflow_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_store=object(),  # type: ignore[arg-type]
    )

    result = AssignmentExecutor(factories).execute(
        _assignment(),
        should_stop=should_stop,
        stop_reason=lambda: "cancellation requested by control plane",
    )

    assert result.status == "cancelled"
    assert result.failure_reason == "cancellation requested by control plane"


def test_workflow_assignment_loads_and_executes_definition() -> None:
    definition = object()
    calls: list[tuple[object, str]] = []

    class FakeWorkflowStore:
        def get_definition(self, definition_id: str):
            return definition if definition_id == "workflow-a" else None

    class FakeWorkflowRunner:
        def run(self, loaded_definition, device_id: str):
            calls.append((loaded_definition, device_id))
            return SimpleNamespace(id="run-a", status="completed")

    factories = ExecutionFactories(
        task_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_runner_factory=lambda: FakeWorkflowRunner(),  # type: ignore[arg-type,return-value]
        workflow_store=FakeWorkflowStore(),  # type: ignore[arg-type]
    )

    result = AssignmentExecutor(factories).execute(
        _assignment(goal=None, workflow_definition_id="workflow-a")
    )

    assert result.status == "done"
    assert calls == [(definition, "device-a")]
    assert result.metadata == {
        "workflow_run_id": "run-a",
        "workflow_status": "completed",
    }


def test_workflow_assignment_maps_cancellation_stop_to_cancelled_status() -> None:
    calls = {"count": 0}

    def should_stop() -> bool:
        calls["count"] += 1
        return calls["count"] > 1

    class FakeWorkflowStore:
        def get_definition(self, definition_id: str):
            return object() if definition_id == "workflow-a" else None

    class FakeWorkflowRunner:
        def run(self, loaded_definition, device_id: str, *, should_stop=None, stop_reason=None):
            assert should_stop is not None and should_stop()
            assert stop_reason is not None
            return SimpleNamespace(
                id="run-a", status="cancelled", failure_reason=stop_reason()
            )

    factories = ExecutionFactories(
        task_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_runner_factory=lambda: FakeWorkflowRunner(),  # type: ignore[arg-type,return-value]
        workflow_store=FakeWorkflowStore(),  # type: ignore[arg-type]
    )

    result = AssignmentExecutor(factories).execute(
        _assignment(goal=None, workflow_definition_id="workflow-a"),
        should_stop=should_stop,
        stop_reason=lambda: "cancellation requested by control plane",
    )

    assert result.status == "cancelled"
    assert result.metadata == {
        "workflow_run_id": "run-a",
        "workflow_status": "cancelled",
    }


def test_unknown_workflow_fails_without_running() -> None:
    class FakeWorkflowStore:
        def get_definition(self, definition_id: str):
            return None

    factories = ExecutionFactories(
        task_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_runner_factory=lambda: object(),  # type: ignore[arg-type,return-value]
        workflow_store=FakeWorkflowStore(),  # type: ignore[arg-type]
    )

    result = AssignmentExecutor(factories).execute(
        _assignment(goal=None, workflow_definition_id="missing")
    )

    assert result.status == "failed"
    assert result.failure_reason == "unknown workflow definition 'missing'"