129 lines
4.4 KiB
Python
129 lines
4.4 KiB
Python
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_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_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'"
|