agentic-mobile-control/tests/test_task_dispatcher.py

"""Unit tests for cloud.dispatch.TaskDispatcher (task 5.6)."""

from __future__ import annotations

from datetime import UTC, datetime
from types import SimpleNamespace

import pytest

from cloud.config import CloudConfig
from cloud.dispatch import (
    Assignment,
    RemoteDispatchNotSupportedError,
    TaskDispatcher,
    UnknownWorkflowDefinitionError,
)
from cloud.pool import DevicePool
from cloud.scheduler import ScheduledTask, TaskConstraints
from cloud.store import CloudStore
from core.models import Task
from workflow.models import (
    PlannedGoalStep,
    WorkflowDefinition,
    WorkflowRun,
    WorkflowStepResult,
)


def _config() -> CloudConfig:
    return CloudConfig(
        sync_interval_seconds=30,
        stale_after_seconds=60,
        max_queue_depth=100,
        default_assignment_strategy="fifo_match",
        api_version_prefix="/v1",
        db_path="cloud/cloud.sqlite3",
    )


class _FakeTaskRunner:
    """A stub TaskRunner that records runs and returns a configured status."""

    def __init__(self, *, status: str = "completed") -> None:
        self._status = status
        self.calls: list[Task] = []

    def run(self, task: Task) -> Task:
        self.calls.append(task)
        task.status = self._status  # type: ignore[assignment]
        if self._status == "completed":
            task.completed_at = datetime.now(UTC)
        elif self._status == "failed":
            task.completed_at = datetime.now(UTC)
            task.failure_reason = "stub failure"
        return task


class _FakeWorkflowStore:
    def __init__(self, definitions: dict[str, WorkflowDefinition] | None = None) -> None:
        self._definitions = definitions or {}

    def get_definition(self, definition_id: str) -> WorkflowDefinition | None:
        return self._definitions.get(definition_id)


class _FakeWorkflowRunner:
    def __init__(
        self,
        *,
        status: str = "completed",
        store: _FakeWorkflowStore | None = None,
    ) -> None:
        self._status = status
        self.store = store or _FakeWorkflowStore()
        self.calls: list[tuple[WorkflowDefinition, str]] = []

    def run(
        self,
        definition: WorkflowDefinition,
        *,
        device_id: str | None = None,
    ) -> WorkflowRun:
        self.calls.append((definition, device_id or ""))
        run = WorkflowRun(
            definition_id=definition.id,
            status=self._status,  # type: ignore[arg-type]
            current_step_id=definition.entry_step_id,
            variables={},
            device_id=device_id,
        )
        return run


def _enqueue_goal_task(store: CloudStore, task_id: str = "task-1") -> str:
    store.enqueue_task(
        ScheduledTask(
            id=task_id,
            goal="open settings",
            workflow_definition_id=None,
            constraints=TaskConstraints(),
            status="assigned",
            created_at=datetime.now(UTC),
        )
    )
    return task_id


def test_local_goal_dispatch_marks_done(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    task_id = _enqueue_goal_task(store)

    runner = _FakeTaskRunner(status="completed")
    dispatcher = TaskDispatcher(
        local_host_id="host-local",
        task_runner_factory=lambda: runner,
        workflow_runner_factory=lambda: _FakeWorkflowRunner(),
        store=store,
    )
    dispatcher.dispatch(
        Assignment(
            task_id=task_id,
            device_id="dev-1",
            host_id="host-local",
            goal="open settings",
            workflow_definition_id=None,
        )
    )

    assert len(runner.calls) == 1
    assert runner.calls[0].device_id == "dev-1"
    task = store.get_task(task_id)
    assert task is not None
    assert task.status == "done"


def test_local_goal_dispatch_marks_failed(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    task_id = _enqueue_goal_task(store)

    runner = _FakeTaskRunner(status="failed")
    dispatcher = TaskDispatcher(
        local_host_id="host-local",
        task_runner_factory=lambda: runner,
        workflow_runner_factory=lambda: _FakeWorkflowRunner(),
        store=store,
    )
    dispatcher.dispatch(
        Assignment(
            task_id=task_id,
            device_id="dev-1",
            host_id="host-local",
            goal="open settings",
            workflow_definition_id=None,
        )
    )

    task = store.get_task(task_id)
    assert task is not None
    assert task.status == "failed"


def _definition() -> WorkflowDefinition:
    return WorkflowDefinition(
        name="linear",
        entry_step_id="first",
        steps=[PlannedGoalStep("first", "do thing")],
    )


def _enqueue_workflow_task(store: CloudStore, definition_id: str, task_id: str = "task-wf") -> str:
    store.enqueue_task(
        ScheduledTask(
            id=task_id,
            goal=None,
            workflow_definition_id=definition_id,
            constraints=TaskConstraints(),
            status="assigned",
            created_at=datetime.now(UTC),
        )
    )
    return task_id


def test_local_workflow_dispatch_runs_definition_and_updates_status(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    definition = _definition()
    wf_store = _FakeWorkflowStore({definition.id: definition})
    runner = _FakeWorkflowRunner(status="completed", store=wf_store)
    task_id = _enqueue_workflow_task(store, definition.id)

    dispatcher = TaskDispatcher(
        local_host_id="host-local",
        task_runner_factory=lambda: _FakeTaskRunner(),
        workflow_runner_factory=lambda: runner,
        store=store,
    )
    dispatcher.dispatch(
        Assignment(
            task_id=task_id,
            device_id="dev-1",
            host_id="host-local",
            goal=None,
            workflow_definition_id=definition.id,
        )
    )

    assert len(runner.calls) == 1
    called_definition, called_device_id = runner.calls[0]
    assert called_definition.id == definition.id
    assert called_device_id == "dev-1"
    task = store.get_task(task_id)
    assert task is not None
    assert task.status == "done"


def test_local_workflow_dispatch_marks_failed(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    definition = _definition()
    wf_store = _FakeWorkflowStore({definition.id: definition})
    runner = _FakeWorkflowRunner(status="failed", store=wf_store)
    task_id = _enqueue_workflow_task(store, definition.id)

    dispatcher = TaskDispatcher(
        local_host_id="host-local",
        task_runner_factory=lambda: _FakeTaskRunner(),
        workflow_runner_factory=lambda: runner,
        store=store,
    )
    dispatcher.dispatch(
        Assignment(
            task_id=task_id,
            device_id="dev-1",
            host_id="host-local",
            goal=None,
            workflow_definition_id=definition.id,
        )
    )

    task = store.get_task(task_id)
    assert task is not None
    assert task.status == "failed"


def test_remote_assignment_raises_and_leaves_assigned(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    task_id = _enqueue_goal_task(store)

    runner = _FakeTaskRunner()
    dispatcher = TaskDispatcher(
        local_host_id="host-local",
        task_runner_factory=lambda: runner,
        workflow_runner_factory=lambda: _FakeWorkflowRunner(),
        store=store,
    )
    with pytest.raises(RemoteDispatchNotSupportedError):
        dispatcher.dispatch(
            Assignment(
                task_id=task_id,
                device_id="dev-remote",
                host_id="host-remote",
                goal="open settings",
                workflow_definition_id=None,
            )
        )

    # The stubbed runner must not have been called.
    assert runner.calls == []
    # Status must remain unchanged from its pre-dispatch value.
    task = store.get_task(task_id)
    assert task is not None
    assert task.status == "assigned"


def test_workflow_dispatch_with_unknown_definition_raises(tmp_path) -> None:
    store = CloudStore(tmp_path / "cloud.sqlite3")
    task_id = _enqueue_workflow_task(store, "missing-def")

    dispatcher = TaskDispatcher(
        local_host_id="host-local",
        task_runner_factory=lambda: _FakeTaskRunner(),
        workflow_runner_factory=lambda: _FakeWorkflowRunner(store=_FakeWorkflowStore({})),
        store=store,
    )
    with pytest.raises(UnknownWorkflowDefinitionError):
        dispatcher.dispatch(
            Assignment(
                task_id=task_id,
                device_id="dev-1",
                host_id="host-local",
                goal=None,
                workflow_definition_id="missing-def",
            )
        )