"""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")