67 lines
2.0 KiB
Python
67 lines
2.0 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
|
|
from cloud.internal_api.models import AssignmentModel
|
|
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",
|
|
workflow_definition_id="workflow-a",
|
|
)
|
|
|
|
|
|
def test_mark_assignment_started_reflected_in_snapshot() -> None:
|
|
ticks = iter([datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC)])
|
|
tracker = AgentStatusTracker(now=lambda: next(ticks))
|
|
|
|
tracker.mark_assignment_started(_assignment())
|
|
|
|
snapshot = tracker.snapshot()
|
|
assert snapshot["current_assignment"] == {
|
|
"task_id": "task-a",
|
|
"device_id": "device-a",
|
|
"goal": "open settings",
|
|
"workflow_definition_id": "workflow-a",
|
|
"started_at": "2026-07-13T09:00:00+00:00",
|
|
}
|
|
|
|
|
|
def test_mark_assignment_finished_clears_current_assignment() -> None:
|
|
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))
|
|
|
|
tracker.mark_assignment_started(_assignment())
|
|
tracker.mark_assignment_finished()
|
|
|
|
assert tracker.snapshot()["current_assignment"] is None
|
|
|
|
|
|
def test_mark_heartbeat_reflected_in_snapshot() -> None:
|
|
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 5, 0, tzinfo=UTC))
|
|
|
|
tracker.mark_heartbeat(ok=True, device_count=3)
|
|
|
|
snapshot = tracker.snapshot()
|
|
assert snapshot["last_heartbeat"] == {
|
|
"ok": True,
|
|
"device_count": 3,
|
|
"at": "2026-07-13T09:05:00+00:00",
|
|
}
|
|
|
|
|
|
def test_snapshot_defaults_to_no_assignment_or_heartbeat() -> None:
|
|
tracker = AgentStatusTracker(now=lambda: datetime(2026, 7, 13, 9, 0, 0, tzinfo=UTC))
|
|
|
|
snapshot = tracker.snapshot()
|
|
|
|
assert snapshot["current_assignment"] is None
|
|
assert snapshot["last_heartbeat"] is None
|