155 lines
4.9 KiB
Python
155 lines
4.9 KiB
Python
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_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())
|