This commit is contained in:
@@ -0,0 +1,66 @@
|
||||
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
|
||||
Reference in New Issue
Block a user