74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
from __future__ import annotations
|
|
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from task_plane.api import TaskPlaneService
|
|
from task_plane.errors import ConflictError, LockedError
|
|
from task_plane.policy import ALL_ACTION_CAPABILITIES, PolicyEngine
|
|
from task_plane.store import TaskPlaneStore
|
|
|
|
|
|
IDENTITY = "spiffe://aios/agent/heartbeat"
|
|
WORKER_ID = "heartbeat-worker"
|
|
|
|
|
|
def claimed_task():
|
|
store = TaskPlaneStore(":memory:")
|
|
service = TaskPlaneService(
|
|
store,
|
|
policy=PolicyEngine(
|
|
execution_enabled=True,
|
|
identity_capabilities={"*": set(ALL_ACTION_CAPABILITIES)},
|
|
),
|
|
)
|
|
service.register_worker(
|
|
{
|
|
"worker_id": WORKER_ID,
|
|
"worker_type": "test",
|
|
"identity": IDENTITY,
|
|
"capabilities": ["code.read"],
|
|
"trust_zone": "test",
|
|
"max_concurrency": 1,
|
|
"heartbeat_interval_seconds": 15,
|
|
"supported_protocols": ["aios-task-v1"],
|
|
}
|
|
)
|
|
task = service.create_task(
|
|
{"task_id": "task-heartbeat", "required_capabilities": ["code.read"]}
|
|
)
|
|
claim = service.claim(
|
|
task["task_id"],
|
|
{
|
|
"worker_id": WORKER_ID,
|
|
"actor_identity": IDENTITY,
|
|
"lease_seconds": 30,
|
|
},
|
|
)
|
|
return store, task, claim
|
|
|
|
|
|
def test_cancelled_task_cannot_renew_a_still_active_claim():
|
|
store, task, claim = claimed_task()
|
|
original_lease = store.get_claim(claim["claim_id"]).lease_until
|
|
current = store.get_task(task["task_id"])
|
|
store.cancel_task_with_event(task["task_id"], current.revision, "owner cancelled")
|
|
|
|
with pytest.raises(ConflictError, match="task is terminal"):
|
|
store.heartbeat_claim(task["task_id"], claim["claim_id"], 60)
|
|
|
|
assert store.get_claim(claim["claim_id"]).lease_until == original_lease
|
|
assert store.get_task(task["task_id"]).status == "CANCELLED"
|
|
|
|
|
|
def test_lease_at_current_time_is_expired_and_cannot_be_renewed():
|
|
store, task, claim = claimed_task()
|
|
boundary = store.get_claim(claim["claim_id"]).lease_until
|
|
|
|
with mock.patch("task_plane.store._now", return_value=boundary):
|
|
with pytest.raises(LockedError, match="lease expired"):
|
|
store.heartbeat_claim(task["task_id"], claim["claim_id"], 30)
|
|
|
|
assert store.get_claim(claim["claim_id"]).lease_until == boundary
|