# -*- coding: utf-8 -*- """Unit tests for ``pineagents.app.task_tracker.TaskTracker``. Covers: - idle/running status before/after a task - external task registration round-trip and idempotency - attach() to non-existent / completed / live runs - attach_or_start() reuses an in-flight run vs. starting a new one - request_stop() cancels and reports running state - detach_subscriber() removes queues and is idempotent - stream_from_queue() yields events and detaches on consumer exit - wait_all_done() returns True when idle, False on timeout - global status counters update via run lifecycle """ # pylint: disable=protected-access,redefined-outer-name,unused-argument from __future__ import annotations import asyncio import json import pytest from pineagents.app.task_tracker import TaskTracker # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- async def _drain(queue: asyncio.Queue, n: int) -> list: """Read up to ``n`` items from ``queue`` with a small timeout.""" items = [] for _ in range(n): items.append(await asyncio.wait_for(queue.get(), timeout=1)) return items def _make_stream(events: list[str]): async def stream(_payload): for ev in events: await asyncio.sleep(0) # cooperate yield ev return stream # --------------------------------------------------------------------------- # get_status / has_active_tasks / list_active_tasks # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_get_status_idle_for_unknown_run_key(): tracker = TaskTracker() assert await tracker.get_status("missing") == "idle" assert await tracker.has_active_tasks() is False assert await tracker.list_active_tasks() == [] @pytest.mark.asyncio async def test_attach_returns_none_for_unknown_run_key(): tracker = TaskTracker() assert await tracker.attach("missing") is None @pytest.mark.asyncio async def test_has_active_tasks_excluding_uses_task_identity(): tracker = TaskTracker() started = asyncio.Event() release = asyncio.Event() producer_sees_other: list[bool] = [] async def producer(_payload): producer_sees_other.append( await tracker.has_active_tasks_excluding( asyncio.current_task(), ), ) started.set() await release.wait() yield "data: done\n\n" queue, _ = await tracker.attach_or_start( "tracked-producer", None, producer, ) await asyncio.wait_for(started.wait(), timeout=1) assert producer_sees_other == [False] assert await tracker.has_active_tasks_excluding( asyncio.current_task(), ) release.set() async for _ in tracker.stream_from_queue(queue, "tracked-producer"): pass # --------------------------------------------------------------------------- # attach_or_start: producer/consumer flow # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_attach_or_start_streams_events_and_marks_completion(): tracker = TaskTracker() events = ["data: a\n\n", "data: b\n\n"] queue, is_new = await tracker.attach_or_start( "run-1", payload=None, stream_fn=_make_stream(events), ) assert is_new is True # Drain the two real events plus the SENTINEL terminator. a = await asyncio.wait_for(queue.get(), timeout=1) b = await asyncio.wait_for(queue.get(), timeout=1) sentinel = await asyncio.wait_for(queue.get(), timeout=1) assert [a, b] == events assert sentinel is None # After completion the tracker cleans up the run. assert await tracker.get_status("run-1") == "idle" assert "run-1" not in tracker._runs @pytest.mark.asyncio async def test_attach_or_start_existing_run_returns_buffer_replay(): tracker = TaskTracker() started = asyncio.Event() release = asyncio.Event() async def slow_stream(_payload): yield "data: first\n\n" started.set() await release.wait() yield "data: second\n\n" queue_a, new_a = await tracker.attach_or_start( "run-2", payload=None, stream_fn=slow_stream, ) assert new_a is True # Wait until the producer has yielded the first event so the buffer # contains something to replay. await asyncio.wait_for(started.wait(), timeout=1) # Yield once more so the broadcast under the lock completes before # the second attach_or_start tries to read the buffer. await asyncio.sleep(0) queue_b, new_b = await tracker.attach_or_start( "run-2", payload=None, stream_fn=_make_stream([]), # must NOT be invoked ) assert new_b is False # queue_b should be pre-filled with the buffered first event. first_b = await asyncio.wait_for(queue_b.get(), timeout=1) assert first_b == "data: first\n\n" # Let the producer finish. release.set() # Both queues see the remaining events and the terminator. rest_a = await _drain(queue_a, 3) # first, second, SENTINEL rest_b = await _drain(queue_b, 2) # second, SENTINEL assert rest_a == ["data: first\n\n", "data: second\n\n", None] assert rest_b == ["data: second\n\n", None] # --------------------------------------------------------------------------- # request_stop: cancellation # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_request_stop_cancels_live_run(): tracker = TaskTracker() started = asyncio.Event() async def long_stream(_payload): started.set() await asyncio.sleep(60) yield "never" await tracker.attach_or_start( "run-cancel", payload=None, stream_fn=long_stream, ) await asyncio.wait_for(started.wait(), timeout=1) assert await tracker.get_status("run-cancel") == "running" stopped = await tracker.request_stop("run-cancel") assert stopped is True # Give the task loop time to process cancellation and cleanup. await asyncio.sleep(0.05) assert await tracker.get_status("run-cancel") == "idle" @pytest.mark.asyncio async def test_request_stop_returns_false_when_no_run(): tracker = TaskTracker() assert await tracker.request_stop("missing") is False # --------------------------------------------------------------------------- # Error path: producer exception broadcasts an error SSE. # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_producer_exception_emits_error_sse(): tracker = TaskTracker() async def boom(_payload): # Make the function an async generator without yielding anything, # so attach_or_start treats it like a real stream that errors out. if False: # pylint: disable=using-constant-test yield raise RuntimeError("kaboom") queue, _ = await tracker.attach_or_start( "run-error", payload=None, stream_fn=boom, ) err = await asyncio.wait_for(queue.get(), timeout=1) sentinel = await asyncio.wait_for(queue.get(), timeout=1) assert err.startswith("data: ") payload = json.loads(err[len("data: ") :].rstrip("\n")) assert payload == {"error": "internal server error"} assert sentinel is None # --------------------------------------------------------------------------- # detach_subscriber: idempotent # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_detach_subscriber_is_idempotent(): tracker = TaskTracker() started = asyncio.Event() release = asyncio.Event() async def gated(_payload): started.set() await release.wait() yield "data: done\n\n" queue, _ = await tracker.attach_or_start( "run-detach", payload=None, stream_fn=gated, ) await asyncio.wait_for(started.wait(), timeout=1) # Detach twice — second call is a no-op. await tracker.detach_subscriber("run-detach", queue) await tracker.detach_subscriber("run-detach", queue) # Detaching a never-registered run also no-ops. await tracker.detach_subscriber("nope", queue) release.set() # Drain to allow producer cleanup. await asyncio.sleep(0.05) # --------------------------------------------------------------------------- # stream_from_queue: consumer detaches on exit # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_stream_from_queue_yields_until_sentinel_and_detaches(): tracker = TaskTracker() events = ["data: 1\n\n", "data: 2\n\n"] queue, _ = await tracker.attach_or_start( "run-stream", payload=None, stream_fn=_make_stream(events), ) collected = [ item async for item in tracker.stream_from_queue(queue, "run-stream") ] assert collected == events # After streaming, run is cleaned up, so detach should be a no-op. assert await tracker.get_status("run-stream") == "idle" # --------------------------------------------------------------------------- # wait_all_done: timeout behaviour # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_wait_all_done_returns_true_when_idle(): tracker = TaskTracker() assert await tracker.wait_all_done(timeout=0.5) is True @pytest.mark.asyncio async def test_wait_all_done_times_out_when_task_runs(): tracker = TaskTracker() release = asyncio.Event() async def producer(_payload): await release.wait() yield "data: done\n\n" queue, _ = await tracker.attach_or_start( "run-long", payload=None, stream_fn=producer, ) try: assert await tracker.wait_all_done(timeout=0.2) is False finally: release.set() async for _ in tracker.stream_from_queue(queue, "run-long"): pass @pytest.mark.asyncio async def test_snapshot_active_tasks_filters_by_owner(): tracker = TaskTracker() owner_a = object() owner_b = object() release = asyncio.Event() async def producer(_payload): await release.wait() yield "data: done\n\n" queue_a, _ = await tracker.attach_or_start( "run-owner-a", None, producer, owner=owner_a, ) queue_b, _ = await tracker.attach_or_start( "run-owner-b", None, producer, owner=owner_b, ) try: snapshot = await tracker.snapshot_active_tasks(owner=owner_a) assert list(snapshot) == ["run-owner-a"] finally: release.set() async for _ in tracker.stream_from_queue(queue_a, "run-owner-a"): pass async for _ in tracker.stream_from_queue(queue_b, "run-owner-b"): pass @pytest.mark.asyncio async def test_wait_tasks_done_ignores_runs_started_after_snapshot(): tracker = TaskTracker() release_old = asyncio.Event() release_new = asyncio.Event() async def old_producer(_payload): await release_old.wait() yield "data: old\n\n" async def new_producer(_payload): await release_new.wait() yield "data: new\n\n" old_queue, _ = await tracker.attach_or_start( "run-old", None, old_producer, ) snapshot = await tracker.snapshot_active_tasks() new_queue, _ = await tracker.attach_or_start( "run-new", None, new_producer, ) release_old.set() assert await tracker.wait_tasks_done( list(snapshot.values()), timeout=1, ) assert await tracker.get_status("run-new") == "running" release_new.set() async for _ in tracker.stream_from_queue(old_queue, "run-old"): pass async for _ in tracker.stream_from_queue(new_queue, "run-new"): pass # --------------------------------------------------------------------------- # Concurrent attach / start safety # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_concurrent_attach_or_start_only_one_producer(): tracker = TaskTracker() invocations = 0 release = asyncio.Event() async def producer(_payload): nonlocal invocations invocations += 1 await release.wait() yield "data: done\n\n" queues = await asyncio.gather( tracker.attach_or_start("run-concurrent", None, producer), tracker.attach_or_start("run-concurrent", None, producer), tracker.attach_or_start("run-concurrent", None, producer), ) new_flags = [is_new for _, is_new in queues] assert new_flags.count(True) == 1 assert invocations == 1 release.set() # Let the producer finish so the test does not leak background tasks. for q, _ in queues: while True: item = await asyncio.wait_for(q.get(), timeout=1) if item is None: break # --------------------------------------------------------------------------- # attach(): replay-end marker for reconnect fast-forward # --------------------------------------------------------------------------- @pytest.mark.asyncio async def test_attach_appends_replay_end_marker_after_buffer(): """Reconnect subscribers get the buffered events, then a ``replay_end`` marker, then live events. The marker lets the client render the replayed part instantly instead of re-animating it.""" tracker = TaskTracker() started = asyncio.Event() release = asyncio.Event() async def slow_stream(_payload): yield "data: first\n\n" started.set() await release.wait() yield "data: second\n\n" queue_a, _ = await tracker.attach_or_start( "run-replay", payload=None, stream_fn=slow_stream, ) await asyncio.wait_for(started.wait(), timeout=1) await asyncio.sleep(0) queue_b = await tracker.attach("run-replay") assert queue_b is not None first = await asyncio.wait_for(queue_b.get(), timeout=1) marker = await asyncio.wait_for(queue_b.get(), timeout=1) assert first == "data: first\n\n" assert marker.startswith("data: ") assert json.loads(marker[len("data: ") :].strip()) == { "type": "replay_end", } release.set() rest_b = await _drain(queue_b, 2) assert rest_b == ["data: second\n\n", None] # The original (non-reconnect) subscriber never sees the marker. rest_a = await _drain(queue_a, 3) assert rest_a == ["data: first\n\n", "data: second\n\n", None]