0cca4963d1
Both the single-run path and the campaign scheduler drove long-lived asyncio tasks through their own duplicated _tasks/_cancel_events dicts and shutdown loops. Collapse them into one deep TaskRegistry module, instantiated as run_registry and campaign_registry. launch() creates the cancel event before the task (so a cancel during startup is never lost), wires done-callback cleanup, and is idempotent per id; this makes runs.py's hard-cancel fallback provably dead, so it is removed. App shutdown now gracefully stops in-flight runs too, not just campaigns.
85 lines
2.3 KiB
Python
85 lines
2.3 KiB
Python
"""Unit tests for TaskRegistry — the durable in-process task registry.
|
|
|
|
Covers the single seam both runs and campaigns share: launch (with an
|
|
injected cancel event), idempotency, cooperative cancel, done-callback
|
|
cleanup, and graceful shutdown_all.
|
|
"""
|
|
|
|
import asyncio
|
|
|
|
from agenteval.task_registry import TaskRegistry
|
|
|
|
|
|
async def test_launch_tracks_task_and_passes_cancel_event():
|
|
reg = TaskRegistry()
|
|
seen: dict = {}
|
|
|
|
async def coro(cancel: asyncio.Event):
|
|
seen["event"] = cancel
|
|
await cancel.wait()
|
|
|
|
task = reg.launch("a", lambda ev: coro(ev))
|
|
assert reg.is_running("a")
|
|
await asyncio.sleep(0) # let the coroutine start and capture the event
|
|
assert isinstance(seen["event"], asyncio.Event)
|
|
|
|
reg.cancel("a")
|
|
await task
|
|
assert not reg.is_running("a") # done-callback popped it
|
|
|
|
|
|
async def test_launch_is_idempotent_for_live_id():
|
|
reg = TaskRegistry()
|
|
|
|
async def coro(cancel: asyncio.Event):
|
|
await cancel.wait()
|
|
|
|
t1 = reg.launch("a", lambda ev: coro(ev))
|
|
t2 = reg.launch("a", lambda ev: coro(ev))
|
|
assert t1 is t2 # same live task, second coro never started
|
|
|
|
reg.cancel("a")
|
|
await t1
|
|
|
|
|
|
async def test_cancel_returns_true_when_live_false_otherwise():
|
|
reg = TaskRegistry()
|
|
|
|
async def coro(cancel: asyncio.Event):
|
|
await cancel.wait()
|
|
|
|
reg.launch("a", lambda ev: coro(ev))
|
|
assert reg.cancel("a") is True # live event signalled
|
|
assert reg.cancel("missing") is False # nothing to signal
|
|
await asyncio.sleep(0) # let the cancelled coro finish
|
|
|
|
|
|
async def test_done_callback_cleans_up_after_natural_completion():
|
|
reg = TaskRegistry()
|
|
|
|
async def quick(cancel: asyncio.Event):
|
|
return "done"
|
|
|
|
task = reg.launch("a", lambda ev: quick(ev))
|
|
await task
|
|
await asyncio.sleep(0) # allow done-callback to run
|
|
assert not reg.is_running("a")
|
|
|
|
|
|
async def test_shutdown_all_cancels_and_awaits_every_task():
|
|
reg = TaskRegistry()
|
|
started = []
|
|
|
|
async def coro(cancel: asyncio.Event):
|
|
started.append(1)
|
|
await cancel.wait()
|
|
|
|
reg.launch("a", lambda ev: coro(ev))
|
|
reg.launch("b", lambda ev: coro(ev))
|
|
await asyncio.sleep(0) # let both start
|
|
assert len(started) == 2
|
|
|
|
await reg.shutdown_all()
|
|
assert not reg.is_running("a")
|
|
assert not reg.is_running("b")
|