From 402cf8dadeec93498de5c67e42fdfa8aa2337898 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Thu, 14 May 2026 09:25:32 -0700 Subject: [PATCH] Port multi-task unit tests from pipecat-subagents Brings over 215 tests across 15 files covering the new multi-task framework: BaseTask / PipelineTask bus lifecycle, job RPC and job groups, the bus message hierarchy and serializers, TaskBus + AsyncQueueBus + RedisBus + PgmqBus (with direct and isolated backends), TaskRegistry, the BusBridgeProcessor, the WebSocket proxy tasks, the LLMTask deferral logic, and the PipelineRunner spawn-and-attach flow. --- .github/workflows/coverage.yaml | 2 + .github/workflows/tests.yaml | 2 + tests/test_base_task.py | 1343 +++++++++++++++++++++++++++++++ tests/test_bridge_processor.py | 267 ++++++ tests/test_bus.py | 274 +++++++ tests/test_bus_network.py | 35 + tests/test_frame_adapters.py | 163 ++++ tests/test_job_group.py | 739 +++++++++++++++++ tests/test_llm_task.py | 285 +++++++ tests/test_messages.py | 97 +++ tests/test_pgmq_backends.py | 266 ++++++ tests/test_pgmq_bus.py | 611 ++++++++++++++ tests/test_redis_bus.py | 249 ++++++ tests/test_registry.py | 170 ++++ tests/test_runner.py | 202 +++++ tests/test_serializers.py | 246 ++++++ tests/test_websocket_proxy.py | 318 ++++++++ 17 files changed, 5269 insertions(+) create mode 100644 tests/test_base_task.py create mode 100644 tests/test_bridge_processor.py create mode 100644 tests/test_bus.py create mode 100644 tests/test_bus_network.py create mode 100644 tests/test_frame_adapters.py create mode 100644 tests/test_job_group.py create mode 100644 tests/test_llm_task.py create mode 100644 tests/test_messages.py create mode 100644 tests/test_pgmq_backends.py create mode 100644 tests/test_pgmq_bus.py create mode 100644 tests/test_redis_bus.py create mode 100644 tests/test_registry.py create mode 100644 tests/test_runner.py create mode 100644 tests/test_serializers.py create mode 100644 tests/test_websocket_proxy.py diff --git a/.github/workflows/coverage.yaml b/.github/workflows/coverage.yaml index 5883afc8a..9386dab9b 100644 --- a/.github/workflows/coverage.yaml +++ b/.github/workflows/coverage.yaml @@ -41,7 +41,9 @@ jobs: --extra google \ --extra langchain \ --extra livekit \ + --extra pgmq \ --extra piper \ + --extra redis \ --extra runner \ --extra sagemaker \ --extra tracing \ diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 04ff89543..7fba47fb7 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -45,7 +45,9 @@ jobs: --extra google \ --extra langchain \ --extra livekit \ + --extra pgmq \ --extra piper \ + --extra redis \ --extra runner \ --extra sagemaker \ --extra tracing \ diff --git a/tests/test_base_task.py b/tests/test_base_task.py new file mode 100644 index 000000000..6552c8d09 --- /dev/null +++ b/tests/test_base_task.py @@ -0,0 +1,1343 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest + +from pipecat.bus import ( + AsyncQueueBus, + BusActivateTaskMessage, + BusAddTaskMessage, + BusCancelMessage, + BusCancelTaskMessage, + BusDeactivateTaskMessage, + BusEndMessage, + BusEndTaskMessage, + BusFrameMessage, + BusJobCancelMessage, + BusJobRequestMessage, + BusJobResponseMessage, + BusJobStreamDataMessage, + BusJobStreamEndMessage, + BusJobStreamStartMessage, + BusJobUpdateMessage, +) +from pipecat.frames.frames import EndFrame, Frame, TextFrame +from pipecat.pipeline.base_task import BaseTask +from pipecat.pipeline.job_context import JobStatus +from pipecat.pipeline.job_decorator import job +from pipecat.pipeline.pipeline import Pipeline +from pipecat.pipeline.runner import PipelineRunner +from pipecat.pipeline.task import PipelineTask +from pipecat.processors.filters.identity_filter import IdentityFilter +from pipecat.processors.frame_processor import FrameDirection, FrameProcessor +from pipecat.registry import TaskRegistry +from pipecat.registry.types import TaskReadyData +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + + +class _FrameGenerator(FrameProcessor): + """Generates a new TextFrame for each input TextFrame.""" + + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + if isinstance(frame, TextFrame): + await self.push_frame(TextFrame(text=f"generated_{frame.text}"), direction) + else: + await self.push_frame(frame, direction) + + +async def create_test_bus(): + """Create an AsyncQueueBus with a TaskManager for testing.""" + bus = AsyncQueueBus() + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + return bus, tm + + +def create_test_registry(): + """Create a registry for testing task lifecycle.""" + return TaskRegistry(runner_name="test-runner") + + +async def register_tasks(registry, *names): + """Pre-register task names so the ready-wait completes immediately.""" + for name in names: + await registry.register(TaskReadyData(task_name=name, runner="test-runner")) + + +def capture_bus(bus): + """Monkey-patch bus.send to capture sent messages in a list.""" + sent = [] + original_send = bus.send + + async def capture_send(message): + sent.append(message) + await original_send(message) + + bus.send = capture_send + return sent + + +def make_stub_pipeline_task(name, *, bridged=None, active=True): + """Create a PipelineTask with an IdentityFilter pipeline.""" + return PipelineTask( + Pipeline([IdentityFilter()]), + name=name, + bridged=bridged, + cancel_on_idle_timeout=False, + ) + + +class TestPipelineTaskLifecycle(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm = await create_test_bus() + self.registry = create_test_registry() + + async def test_task_starts_inactive_by_default(self): + """Bridged task is inactive by default.""" + task = make_stub_pipeline_task("test", bridged=()) + task._active = False + task._pending_activation = False + self.assertFalse(task.active) + + async def test_handoff_via_bus_message_after_pipeline_start(self): + """Task activates when BusActivateTaskMessage received and pipeline started.""" + task = make_stub_pipeline_task("test", bridged=()) + task._active = False + task._pending_activation = False + task.attach(registry=self.registry, bus=self.bus) + + handoff_done = asyncio.Event() + handoff_args_received = [] + + @task.event_handler("on_activated") + async def on_handoff(t, args): + handoff_args_received.append(args) + handoff_done.set() + + args = {"messages": ["hello"]} + + async def handoff_after_start(): + await asyncio.sleep(0.05) + await self.bus.send(BusActivateTaskMessage(source="other", target="test", args=args)) + await asyncio.wait_for(handoff_done.wait(), timeout=2.0) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), handoff_after_start()) + + self.assertTrue(task.active) + self.assertEqual(len(handoff_args_received), 1) + self.assertEqual(handoff_args_received[0], args) + + async def test_active_true_fires_on_activated(self): + """active=True fires on_activated after pipeline starts.""" + task = make_stub_pipeline_task("test", bridged=()) + + activated = asyncio.Event() + + @task.event_handler("on_activated") + async def on_activated(t, args): + activated.set() + + async def wait_and_end(): + await asyncio.wait_for(activated.wait(), timeout=2.0) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), wait_and_end()) + + self.assertTrue(task.active) + self.assertTrue(activated.is_set()) + + async def test_activation_args_property_set_and_cleared(self): + """activation_args returns the latest args while active and is cleared on deactivate.""" + task = make_stub_pipeline_task("test", bridged=()) + task._active = False + task._pending_activation = False + + activated = asyncio.Event() + deactivated = asyncio.Event() + + @task.event_handler("on_activated") + async def _on_activated(t, args): + activated.set() + + @task.event_handler("on_deactivated") + async def _on_deactivated(t): + deactivated.set() + + args = {"messages": ["hello"]} + observed_while_active = {} + + async def drive(): + await asyncio.sleep(0.05) + await self.bus.send(BusActivateTaskMessage(source="other", target="test", args=args)) + await asyncio.wait_for(activated.wait(), timeout=2.0) + observed_while_active["args"] = task.activation_args + observed_while_active["active"] = task.active + await self.bus.send(BusDeactivateTaskMessage(source="other", target="test")) + await asyncio.wait_for(deactivated.wait(), timeout=2.0) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), drive()) + + self.assertTrue(observed_while_active["active"]) + self.assertEqual(observed_while_active["args"], args) + self.assertFalse(task.active) + self.assertIsNone(task.activation_args) + + async def test_activation_args_none_before_activation(self): + """activation_args is None before any activation has occurred.""" + task = make_stub_pipeline_task("test", bridged=()) + task._active = False + task._pending_activation = False + self.assertIsNone(task.activation_args) + + async def test_handoff_to_sends_activate_and_deactivates(self): + """handoff_to() sends BusDeactivateTaskMessage and BusActivateTaskMessage.""" + sent = capture_bus(self.bus) + + task = make_stub_pipeline_task("task_a", bridged=()) + task.attach(registry=self.registry, bus=self.bus) + + await task.handoff_to("task_b") + + deactivate_msgs = [m for m in sent if isinstance(m, BusDeactivateTaskMessage)] + self.assertEqual(len(deactivate_msgs), 1) + self.assertEqual(deactivate_msgs[0].target, "task_a") + activate_msgs = [m for m in sent if isinstance(m, BusActivateTaskMessage)] + self.assertEqual(len(activate_msgs), 1) + self.assertEqual(activate_msgs[0].target, "task_b") + + async def test_end_without_parent_sends_bus_end_message(self): + """end() with no parent sends BusEndMessage.""" + sent = capture_bus(self.bus) + + task = BaseTask("task_a") + task.attach(registry=self.registry, bus=self.bus) + await task.end(reason="done") + + end_msgs = [m for m in sent if isinstance(m, BusEndMessage)] + self.assertEqual(len(end_msgs), 1) + self.assertEqual(end_msgs[0].source, "task_a") + self.assertEqual(end_msgs[0].reason, "done") + + async def test_end_with_parent_sends_bus_end_message(self): + """end() with parent still sends BusEndMessage (runner handles it).""" + sent = capture_bus(self.bus) + + parent = BaseTask("parent_task") + parent.attach(registry=self.registry, bus=self.bus) + task = BaseTask("child") + task.attach(registry=self.registry, bus=self.bus) + await parent.add_task(task) + await task.end(reason="goodbye") + + end_msgs = [m for m in sent if isinstance(m, BusEndMessage)] + self.assertEqual(len(end_msgs), 1) + self.assertEqual(end_msgs[0].source, "child") + self.assertEqual(end_msgs[0].reason, "goodbye") + + async def test_cancel_sends_bus_cancel_message(self): + """cancel() sends BusCancelMessage.""" + sent = capture_bus(self.bus) + + task = BaseTask("task_a") + task.attach(registry=self.registry, bus=self.bus) + await task.cancel() + + cancel_msgs = [m for m in sent if isinstance(m, BusCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].source, "task_a") + + async def test_add_task_sends_bus_add_task_message(self): + """add_task() sends BusAddTaskMessage.""" + sent = capture_bus(self.bus) + + task = BaseTask("task_a") + task.attach(registry=self.registry, bus=self.bus) + new_task = BaseTask("task_b") + await task.add_task(new_task) + + add_msgs = [m for m in sent if isinstance(m, BusAddTaskMessage)] + self.assertEqual(len(add_msgs), 1) + self.assertIs(add_msgs[0].task, new_task) + + async def test_started_at_none_before_pipeline_starts(self): + """started_at is None before the pipeline has started.""" + task = make_stub_pipeline_task("test") + self.assertIsNone(task.started_at) + + async def test_started_at_set_after_pipeline_starts(self): + """started_at becomes set once the pipeline starts.""" + task = make_stub_pipeline_task("test") + + ready_event = asyncio.Event() + + @task.event_handler("on_pipeline_started") + async def _on_started(t, frame): + ready_event.set() + + async def wait_and_end(): + await asyncio.wait_for(ready_event.wait(), timeout=2.0) + self.assertIsNotNone(task.started_at) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), wait_and_end()) + + async def test_on_pipeline_started_event(self): + """on_pipeline_started fires after pipeline starts.""" + task = make_stub_pipeline_task("test") + + started = asyncio.Event() + + @task.event_handler("on_pipeline_started") + async def on_started(t, frame): + started.set() + + async def wait_and_end(): + await asyncio.wait_for(started.wait(), timeout=2.0) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), wait_and_end()) + + self.assertTrue(started.is_set()) + + async def test_on_pipeline_finished_event(self): + """on_pipeline_finished fires after pipeline finishes.""" + task = make_stub_pipeline_task("test") + + finished_fired = asyncio.Event() + + @task.event_handler("on_pipeline_finished") + async def on_finished(t, frame): + finished_fired.set() + + async def end_pipeline(): + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), end_pipeline()) + + self.assertTrue(finished_fired.is_set()) + + async def test_handoff_deactivates(self): + """handoff_to() sends a deactivate message for the calling task.""" + sent = capture_bus(self.bus) + task = make_stub_pipeline_task("test", bridged=()) + task.attach(registry=self.registry, bus=self.bus) + + self.assertTrue(task.active) + await task.handoff_to("other") + deactivate_msgs = [m for m in sent if isinstance(m, BusDeactivateTaskMessage)] + self.assertEqual(len(deactivate_msgs), 1) + self.assertEqual(deactivate_msgs[0].target, "test") + + async def test_bus_end_task_message_ends_pipeline(self): + """BusEndTaskMessage causes the pipeline to end gracefully.""" + task = make_stub_pipeline_task("test") + + finished = asyncio.Event() + + @task.event_handler("on_pipeline_finished") + async def on_finished(t, frame): + if isinstance(frame, EndFrame): + finished.set() + + async def send_end_message(): + await asyncio.sleep(0.05) + await self.bus.send( + BusEndTaskMessage(source="runner", target="test", reason="shutdown") + ) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), send_end_message()) + + self.assertTrue(finished.is_set()) + + async def test_bus_cancel_task_message_cancels_pipeline(self): + """BusCancelTaskMessage cancels the pipeline task.""" + task = make_stub_pipeline_task("test") + + async def send_cancel_message(): + await asyncio.sleep(0.05) + await self.bus.send(BusCancelTaskMessage(source="runner", target="test")) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + try: + await asyncio.gather(runner.run(task), send_cancel_message()) + except asyncio.CancelledError: + pass + + self.assertTrue(task.has_finished()) + + async def test_queue_frame(self): + """queue_frame injects a frame into the pipeline.""" + task = make_stub_pipeline_task("test") + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def push_frames(): + await asyncio.sleep(0.05) + await task.queue_frame(TextFrame(text="injected")) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), push_frames()) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].text, "injected") + + async def test_queue_frames(self): + """queue_frames injects multiple frames into the pipeline.""" + task = make_stub_pipeline_task("test") + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def push_frames(): + await asyncio.sleep(0.05) + await task.queue_frames([TextFrame(text="a"), TextFrame(text="b")]) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), push_frames()) + + self.assertEqual(len(received), 2) + self.assertEqual(received[0].text, "a") + self.assertEqual(received[1].text, "b") + + async def test_self_handoff(self): + """A task can handoff to itself via handoff_to(self.name).""" + task = make_stub_pipeline_task("test", bridged=()) + + handoff_done = asyncio.Event() + + @task.event_handler("on_activated") + async def on_handoff(t, args): + handoff_done.set() + + async def self_handoff(): + # Wait for first activation (from active=True) + await asyncio.sleep(0.05) + handoff_done.clear() + await task.handoff_to("test") + await asyncio.wait_for(handoff_done.wait(), timeout=2.0) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), self_handoff()) + + self.assertTrue(task.active) + + async def test_add_task_tracks_children(self): + """add_task() populates children list and sets parent.""" + parent = BaseTask("parent") + parent.attach(registry=self.registry, bus=self.bus) + child_a = BaseTask("child_a") + child_b = BaseTask("child_b") + + await parent.add_task(child_a) + await parent.add_task(child_b) + + self.assertEqual(len(parent.children), 2) + self.assertIs(parent.children[0], child_a) + self.assertIs(parent.children[1], child_b) + + async def test_end_propagates_to_children(self): + """BusEndTaskMessage on parent sends end to each child.""" + sent = capture_bus(self.bus) + + parent = BaseTask("parent") + parent.attach(registry=self.registry, bus=self.bus) + child_a = BaseTask("child_a") + child_b = BaseTask("child_b") + await parent.add_task(child_a) + await parent.add_task(child_b) + + # Pre-set children as finished so gather returns immediately + child_a._finished_event.set() + child_b._finished_event.set() + + await parent.on_bus_message( + BusEndTaskMessage(source="runner", target="parent", reason="shutdown") + ) + + end_msgs = [m for m in sent if isinstance(m, BusEndTaskMessage)] + targets = {m.target for m in end_msgs} + self.assertIn("child_a", targets) + self.assertIn("child_b", targets) + + async def test_end_waits_for_children(self): + """Parent waits for children to finish before completing _handle_task_end.""" + parent = BaseTask("parent") + parent.attach(registry=self.registry, bus=self.bus) + child = BaseTask("child") + await parent.add_task(child) + + order = [] + + async def delayed_child_finish(): + await asyncio.sleep(0.1) + order.append("child_finished") + child._finished_event.set() + + async def send_end(): + await asyncio.sleep(0.05) + await parent.on_bus_message( + BusEndTaskMessage(source="runner", target="parent", reason="shutdown") + ) + order.append("parent_end_returned") + + await asyncio.gather(send_end(), delayed_child_finish()) + + # Child must finish before parent's on_bus_message returns + self.assertEqual(order, ["child_finished", "parent_end_returned"]) + + async def test_cancel_propagates_to_children(self): + """BusCancelTaskMessage on parent sends cancel to each child.""" + sent = capture_bus(self.bus) + + parent = BaseTask("parent") + parent.attach(registry=self.registry, bus=self.bus) + child_a = BaseTask("child_a") + child_b = BaseTask("child_b") + await parent.add_task(child_a) + await parent.add_task(child_b) + + await parent.on_bus_message( + BusCancelTaskMessage(source="runner", target="parent", reason="abort") + ) + + cancel_msgs = [m for m in sent if isinstance(m, BusCancelTaskMessage)] + targets = {m.target for m in cancel_msgs} + self.assertIn("child_a", targets) + self.assertIn("child_b", targets) + + +def make_generating_task(name, *, bridged=None): + """Create a PipelineTask whose pipeline generates frames.""" + return PipelineTask( + Pipeline([_FrameGenerator()]), + name=name, + bridged=bridged, + cancel_on_idle_timeout=False, + ) + + +class TestEdgeToBus(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm = await create_test_bus() + self.registry = create_test_registry() + + async def test_generated_frames_reach_bus(self): + """Pipeline-generated frames are broadcast to the bus.""" + sent = capture_bus(self.bus) + + task = make_generating_task("task", bridged=()) + + async def push_frames(): + await asyncio.sleep(0.05) + await task.queue_frame(TextFrame(text="input")) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), push_frames()) + + bus_frame_msgs = [m for m in sent if isinstance(m, BusFrameMessage)] + text_msgs = [m for m in bus_frame_msgs if isinstance(m.frame, TextFrame)] + generated = [m for m in text_msgs if m.frame.text == "generated_input"] + self.assertEqual(len(generated), 1) + self.assertEqual(generated[0].source, "task") + + async def test_bus_frames_not_rebroadcast_by_same_task(self): + """Frames from the bus with source==self are ignored by edge processors.""" + sent = capture_bus(self.bus) + + task = make_stub_pipeline_task("task", bridged=()) + + async def inject_frame(): + await asyncio.sleep(0.05) + # Send a frame from "other" — edge source accepts it (downstream, source != task) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="from_bus"), + direction=FrameDirection.DOWNSTREAM, + ) + ) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), inject_frame()) + + # The frame passes through the identity pipeline and reaches + # EdgeSink, which re-broadcasts with source="task". That's + # expected. But it must NOT create a loop — EdgeSource ignores + # it because source == "task". + # Filter to TextFrame only to ignore metrics frames. + bus_frame_msgs = [ + m for m in sent if isinstance(m, BusFrameMessage) and isinstance(m.frame, TextFrame) + ] + from_task = [m for m in bus_frame_msgs if m.source == "task"] + from_other = [m for m in bus_frame_msgs if m.source == "other"] + # One re-broadcast from task (EdgeSink), one original from other + self.assertEqual(len(from_other), 1) + self.assertEqual(len(from_task), 1) + # No infinite loop — total is exactly 2 + self.assertEqual(len(bus_frame_msgs), 2) + + async def test_unbridged_task_no_edge_sinks(self): + """An unbridged PipelineTask does not tee frames to the bus.""" + sent = capture_bus(self.bus) + + task = make_stub_pipeline_task("root", bridged=None) + + async def push_frames(): + await asyncio.sleep(0.05) + await task.queue_frame(TextFrame(text="root_frame")) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), push_frames()) + + bus_frame_msgs = [m for m in sent if isinstance(m, BusFrameMessage)] + self.assertEqual(len(bus_frame_msgs), 0) + + async def test_bus_frame_enters_task_pipeline(self): + """Bus frame messages enter the pipeline via edge source processor.""" + task = make_stub_pipeline_task("task", bridged=()) + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def inject_frame(): + await asyncio.sleep(0.05) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="from_bus"), + direction=FrameDirection.DOWNSTREAM, + ) + ) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), inject_frame()) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].text, "from_bus") + + async def test_direction_preserved_in_bus_frame(self): + """Direction is preserved when generated frames are sent to the bus.""" + sent = capture_bus(self.bus) + + task = make_generating_task("task", bridged=()) + + async def push_frames(): + await asyncio.sleep(0.05) + await task.queue_frame(TextFrame(text="hello")) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), push_frames()) + + bus_frame_msgs = [m for m in sent if isinstance(m, BusFrameMessage)] + text_msgs = [m for m in bus_frame_msgs if isinstance(m.frame, TextFrame)] + generated = [m for m in text_msgs if m.frame.text == "generated_hello"] + self.assertEqual(len(generated), 1) + self.assertEqual(generated[0].direction, FrameDirection.DOWNSTREAM) + + async def test_bridged_task_accepts_matching_bridge(self): + """Bridged task with named bridge accepts frames from that bridge.""" + task = make_stub_pipeline_task("task", bridged=("voice",)) + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def inject_frame(): + await asyncio.sleep(0.05) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="voice_frame"), + direction=FrameDirection.DOWNSTREAM, + bridge="voice", + ) + ) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), inject_frame()) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].text, "voice_frame") + + async def test_bridged_task_rejects_non_matching_bridge(self): + """Bridged task with named bridge rejects frames from other bridges.""" + task = make_stub_pipeline_task("task", bridged=("voice",)) + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def inject_frame(): + await asyncio.sleep(0.05) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="video_frame"), + direction=FrameDirection.DOWNSTREAM, + bridge="video", + ) + ) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), inject_frame()) + + self.assertEqual(len(received), 0) + + async def test_bridged_task_empty_tuple_accepts_all(self): + """Bridged task with empty tuple accepts frames from any bridge.""" + task = make_stub_pipeline_task("task", bridged=()) + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def inject_frames(): + await asyncio.sleep(0.05) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="voice"), + direction=FrameDirection.DOWNSTREAM, + bridge="voice", + ) + ) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="video"), + direction=FrameDirection.DOWNSTREAM, + bridge="video", + ) + ) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="none"), + direction=FrameDirection.DOWNSTREAM, + ) + ) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), inject_frames()) + + self.assertEqual(len(received), 3) + + async def test_bridged_task_multiple_bridges(self): + """Bridged task with multiple bridge names accepts from all listed.""" + task = make_stub_pipeline_task("task", bridged=("voice", "video")) + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(t, frame): + received.append(frame) + + async def inject_frames(): + await asyncio.sleep(0.05) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="voice"), + direction=FrameDirection.DOWNSTREAM, + bridge="voice", + ) + ) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="video"), + direction=FrameDirection.DOWNSTREAM, + bridge="video", + ) + ) + await self.bus.send( + BusFrameMessage( + source="other", + frame=TextFrame(text="other"), + direction=FrameDirection.DOWNSTREAM, + bridge="other", + ) + ) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), inject_frames()) + + texts = sorted([r.text for r in received]) + self.assertEqual(texts, ["video", "voice"]) + + async def test_not_bridged_task_ignores_bridge(self): + """Non-bridged task (bridged=None) has no edge processors.""" + sent = capture_bus(self.bus) + + task = make_stub_pipeline_task("root", bridged=None) + + async def push_frames(): + await asyncio.sleep(0.05) + await task.queue_frame(TextFrame(text="test")) + await asyncio.sleep(0.05) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner(bus=self.bus, handle_sigint=False) + await asyncio.gather(runner.run(task), push_frames()) + + bus_frame_msgs = [m for m in sent if isinstance(m, BusFrameMessage)] + self.assertEqual(len(bus_frame_msgs), 0) + + +class TestJobLifecycle(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm = await create_test_bus() + self.registry = create_test_registry() + + async def _attach(self, task): + task.attach(registry=self.registry, bus=self.bus) + await task.setup(self.tm) + return task + + async def test_request_job_sends_request(self): + """request_job() sends BusJobRequestMessage to each task.""" + sent = capture_bus(self.bus) + + parent = await self._attach(BaseTask("parent")) + await self.registry.register(TaskReadyData(task_name="worker", runner="test-runner")) + + job_id = await parent.request_job("worker", payload={"key": "val"}) + + request_msgs = [m for m in sent if isinstance(m, BusJobRequestMessage)] + self.assertEqual(len(request_msgs), 1) + self.assertEqual(request_msgs[0].job_id, job_id) + self.assertEqual(request_msgs[0].target, "worker") + self.assertEqual(request_msgs[0].payload, {"key": "val"}) + + async def test_request_job_group_multiple_tasks(self): + """request_job_group() with multiple tasks sends messages for each.""" + sent = capture_bus(self.bus) + + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "w1", "w2") + + job_id = await parent.request_job_group("w1", "w2", payload={"work": True}) + + request_msgs = [m for m in sent if isinstance(m, BusJobRequestMessage)] + self.assertEqual(len(request_msgs), 2) + targets = {m.target for m in request_msgs} + self.assertEqual(targets, {"w1", "w2"}) + for m in request_msgs: + self.assertEqual(m.job_id, job_id) + + async def test_on_job_request_called(self): + """BusJobRequestMessage triggers on_job_request with the message.""" + task = await self._attach(BaseTask("worker")) + + received = [] + + @task.event_handler("on_job_request") + async def on_request(t, message): + received.append(message) + + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1", payload={"x": 1}) + ) + await asyncio.sleep(0) # let async event handlers run + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].job_id, "t1") + self.assertEqual(received[0].source, "parent") + self.assertEqual(received[0].payload, {"x": 1}) + self.assertIn("t1", task.active_jobs) + + async def test_send_job_response(self): + """send_job_response() sends BusJobResponseMessage to requester.""" + sent = capture_bus(self.bus) + + task = await self._attach(BaseTask("worker")) + # Simulate receiving a job request + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + + await task.send_job_response("t1", {"result": 42}) + + response_msgs = [m for m in sent if isinstance(m, BusJobResponseMessage)] + self.assertEqual(len(response_msgs), 1) + self.assertEqual(response_msgs[0].target, "parent") + self.assertEqual(response_msgs[0].job_id, "t1") + self.assertEqual(response_msgs[0].response, {"result": 42}) + self.assertEqual(response_msgs[0].status, JobStatus.COMPLETED) + + async def test_send_job_update(self): + """send_job_update() sends BusJobUpdateMessage to requester.""" + sent = capture_bus(self.bus) + + task = await self._attach(BaseTask("worker")) + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + + await task.send_job_update("t1", {"progress": 50}) + + update_msgs = [m for m in sent if isinstance(m, BusJobUpdateMessage)] + self.assertEqual(len(update_msgs), 1) + self.assertEqual(update_msgs[0].target, "parent") + self.assertEqual(update_msgs[0].job_id, "t1") + self.assertEqual(update_msgs[0].update, {"progress": 50}) + + async def test_on_job_completed_fires_when_all_respond(self): + """on_job_completed fires when all tasks in a job group respond.""" + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "w1", "w2") + + completed = [] + + @parent.event_handler("on_job_completed") + async def on_completed(t, result): + completed.append(result) + + job_id = await parent.request_job_group("w1", "w2") + + # First response — should not trigger on_job_completed + await parent.on_bus_message( + BusJobResponseMessage( + source="w1", + target="parent", + job_id=job_id, + status=JobStatus.COMPLETED, + response={"a": 1}, + ) + ) + await asyncio.sleep(0) # let async event handlers run + self.assertEqual(len(completed), 0) + + # Second response — should trigger on_job_completed + await parent.on_bus_message( + BusJobResponseMessage( + source="w2", + target="parent", + job_id=job_id, + status=JobStatus.COMPLETED, + response={"b": 2}, + ) + ) + await asyncio.sleep(0) # let async event handlers run + self.assertEqual(len(completed), 1) + self.assertEqual(completed[0].job_id, job_id) + self.assertEqual(completed[0].responses, {"w1": {"a": 1}, "w2": {"b": 2}}) + + async def test_cancel_job_group_sends_cancel_to_all_tasks(self): + """cancel_job_group() sends BusJobCancelMessage to all tasks in group.""" + sent = capture_bus(self.bus) + + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "w1", "w2") + + job_id = await parent.request_job_group("w1", "w2") + sent.clear() + + await parent.cancel_job_group(job_id, reason="no longer needed") + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 2) + targets = {m.target for m in cancel_msgs} + self.assertEqual(targets, {"w1", "w2"}) + for m in cancel_msgs: + self.assertEqual(m.job_id, job_id) + self.assertEqual(m.reason, "no longer needed") + + async def test_send_job_response_raises_without_active_job(self): + """send_job_response raises RuntimeError when job_id is unknown.""" + task = await self._attach(BaseTask("worker")) + + with self.assertRaises(RuntimeError): + await task.send_job_response("unknown", {"result": 1}) + + async def test_send_job_update_raises_without_active_job(self): + """send_job_update raises RuntimeError when job_id is unknown.""" + task = await self._attach(BaseTask("worker")) + + with self.assertRaises(RuntimeError): + await task.send_job_update("unknown", {"progress": 50}) + + async def test_cancel_auto_sends_cancelled_response(self): + """BusJobCancelMessage auto-sends a cancelled response and clears state.""" + task = await self._attach(BaseTask("worker")) + + # Set up job state + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + self.assertIn("t1", task.active_jobs) + + # Cancel should auto-send response (via send_job_response) and clear state + await task.on_bus_message( + BusJobCancelMessage(source="parent", target="worker", job_id="t1") + ) + # Yield to let the cancel coroutine complete its work. + await asyncio.sleep(0.05) + self.assertNotIn("t1", task.active_jobs) + + async def test_on_job_cancelled_fires(self): + """BusJobCancelMessage triggers on_job_cancelled with the message.""" + task = await self._attach(BaseTask("worker")) + + received = [] + + @task.event_handler("on_job_cancelled") + async def on_cancelled(t, message): + received.append(message) + + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + await task.on_bus_message( + BusJobCancelMessage( + source="parent", target="worker", job_id="t1", reason="no longer needed" + ) + ) + await asyncio.sleep(0.05) + + self.assertEqual(len(received), 1) + self.assertEqual(received[0].job_id, "t1") + self.assertEqual(received[0].reason, "no longer needed") + + async def test_send_job_stream_start(self): + """send_job_stream_start() sends BusJobStreamStartMessage to requester.""" + sent = capture_bus(self.bus) + + task = await self._attach(BaseTask("worker")) + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + + await task.send_job_stream_start("t1", {"content_type": "text"}) + + msgs = [m for m in sent if isinstance(m, BusJobStreamStartMessage)] + self.assertEqual(len(msgs), 1) + self.assertEqual(msgs[0].target, "parent") + self.assertEqual(msgs[0].job_id, "t1") + self.assertEqual(msgs[0].data, {"content_type": "text"}) + + async def test_send_job_stream_data(self): + """send_job_stream_data() sends BusJobStreamDataMessage to requester.""" + sent = capture_bus(self.bus) + + task = await self._attach(BaseTask("worker")) + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + + await task.send_job_stream_data("t1", {"text": "hello "}) + + msgs = [m for m in sent if isinstance(m, BusJobStreamDataMessage)] + self.assertEqual(len(msgs), 1) + self.assertEqual(msgs[0].target, "parent") + self.assertEqual(msgs[0].job_id, "t1") + self.assertEqual(msgs[0].data, {"text": "hello "}) + + async def test_send_job_stream_end(self): + """send_job_stream_end() sends BusJobStreamEndMessage to requester.""" + sent = capture_bus(self.bus) + + task = await self._attach(BaseTask("worker")) + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="t1") + ) + + await task.send_job_stream_end("t1", {"final": True}) + + msgs = [m for m in sent if isinstance(m, BusJobStreamEndMessage)] + self.assertEqual(len(msgs), 1) + self.assertEqual(msgs[0].target, "parent") + self.assertEqual(msgs[0].job_id, "t1") + self.assertEqual(msgs[0].data, {"final": True}) + + async def test_stream_end_triggers_on_job_completed(self): + """BusJobStreamEndMessage triggers group completion like BusJobResponseMessage.""" + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "w1", "w2") + + completed = [] + + @parent.event_handler("on_job_completed") + async def on_completed(t, result): + completed.append(result) + + job_id = await parent.request_job_group("w1", "w2") + + # First task ends stream — should not trigger on_job_completed + await parent.on_bus_message( + BusJobStreamEndMessage( + source="w1", target="parent", job_id=job_id, data={"result": "a"} + ) + ) + await asyncio.sleep(0) + self.assertEqual(len(completed), 0) + + # Second task ends stream — should trigger on_job_completed + await parent.on_bus_message( + BusJobStreamEndMessage( + source="w2", target="parent", job_id=job_id, data={"result": "b"} + ) + ) + await asyncio.sleep(0) + self.assertEqual(len(completed), 1) + self.assertEqual(completed[0].job_id, job_id) + self.assertEqual(completed[0].responses, {"w1": {"result": "a"}, "w2": {"result": "b"}}) + + async def test_send_job_stream_raises_without_active_job(self): + """All stream helpers raise RuntimeError when job_id is unknown.""" + task = await self._attach(BaseTask("worker")) + + with self.assertRaises(RuntimeError): + await task.send_job_stream_start("unknown") + + with self.assertRaises(RuntimeError): + await task.send_job_stream_data("unknown") + + with self.assertRaises(RuntimeError): + await task.send_job_stream_end("unknown") + + async def test_request_job_timeout_cancels_job(self): + """Short timeout triggers BusJobCancelMessage with reason 'timeout'.""" + sent = capture_bus(self.bus) + + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "worker") + + job_id = await parent.request_job("worker", timeout=0.05) + + # Wait for timeout to fire + await asyncio.sleep(0.1) + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].job_id, job_id) + self.assertEqual(cancel_msgs[0].reason, "timeout") + + # Clean up remaining tasks + await parent.cleanup() + + async def test_request_job_timeout_cancelled_on_completion(self): + """Responding before timeout prevents cancel from being sent.""" + sent = capture_bus(self.bus) + + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "worker") + + job_id = await parent.request_job("worker", timeout=0.5) + + # Let the timeout task start before responding + await asyncio.sleep(0) + + # Respond before timeout fires + await parent.on_bus_message( + BusJobResponseMessage( + source="worker", + target="parent", + job_id=job_id, + status=JobStatus.COMPLETED, + response={"ok": True}, + ) + ) + + # Wait past what would have been the timeout + await asyncio.sleep(0.1) + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 0) + + # Clean up remaining tasks + await parent.cleanup() + + async def test_request_job_no_timeout_by_default(self): + """timeout_task is None when no timeout is given.""" + parent = await self._attach(BaseTask("parent")) + await register_tasks(self.registry, "worker") + + job_id = await parent.request_job("worker") + + group = parent.job_groups[job_id] + self.assertIsNone(group.timeout_task) + + +class TestJobDecorator(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm = await create_test_bus() + self.registry = create_test_registry() + + async def _attach(self, task): + task.attach(registry=self.registry, bus=self.bus) + await task.setup(self.tm) + return task + + async def _wait_until(self, predicate, timeout=1.0): + deadline = asyncio.get_event_loop().time() + timeout + while not predicate(): + if asyncio.get_event_loop().time() > deadline: + raise AssertionError("condition not met within timeout") + await asyncio.sleep(0.01) + + async def test_sequential_runs_one_at_a_time_in_order(self): + """sequential=True serializes same-name requests in FIFO order.""" + running = 0 + max_running = 0 + completion_order: list[str] = [] + gates: dict[str, asyncio.Event] = {} + + class WorkerTask(BaseTask): + @job(name="work", sequential=True) + async def on_work(self, message): + nonlocal running, max_running + running += 1 + max_running = max(max_running, running) + await gates[message.job_id].wait() + completion_order.append(message.job_id) + running -= 1 + + task = await self._attach(WorkerTask("worker")) + + for tid in ("t1", "t2", "t3"): + gates[tid] = asyncio.Event() + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id=tid, job_name="work") + ) + + # Let the first handler enter the locked region. + await self._wait_until(lambda: running == 1) + + # Release in order; verify completion in order. + for tid in ("t1", "t2", "t3"): + gates[tid].set() + await self._wait_until(lambda tid=tid: tid in completion_order) + + self.assertEqual(completion_order, ["t1", "t2", "t3"]) + self.assertEqual(max_running, 1) + + async def test_non_sequential_runs_concurrently(self): + """Without sequential=True, same-name handlers run concurrently.""" + running = 0 + max_running = 0 + release = asyncio.Event() + + class WorkerTask(BaseTask): + @job(name="work") + async def on_work(self, message): + nonlocal running, max_running + running += 1 + max_running = max(max_running, running) + await release.wait() + running -= 1 + + task = await self._attach(WorkerTask("worker")) + + for tid in ("t1", "t2"): + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id=tid, job_name="work") + ) + + await self._wait_until(lambda: running == 2) + self.assertEqual(max_running, 2) + release.set() + await self._wait_until(lambda: running == 0) + + async def test_sequential_locks_are_per_name(self): + """Sequential lock is per job name; different names do not block each other.""" + running = 0 + max_running = 0 + release = asyncio.Event() + + class WorkerTask(BaseTask): + @job(name="a", sequential=True) + async def on_a(self, message): + nonlocal running, max_running + running += 1 + max_running = max(max_running, running) + await release.wait() + running -= 1 + + @job(name="b", sequential=True) + async def on_b(self, message): + nonlocal running, max_running + running += 1 + max_running = max(max_running, running) + await release.wait() + running -= 1 + + task = await self._attach(WorkerTask("worker")) + + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="ta", job_name="a") + ) + await task.on_bus_message( + BusJobRequestMessage(source="parent", target="worker", job_id="tb", job_name="b") + ) + + await self._wait_until(lambda: running == 2) + self.assertEqual(max_running, 2) + release.set() + await self._wait_until(lambda: running == 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bridge_processor.py b/tests/test_bridge_processor.py new file mode 100644 index 000000000..92f611cb5 --- /dev/null +++ b/tests/test_bridge_processor.py @@ -0,0 +1,267 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest + +from pipecat.bus import AsyncQueueBus, BusBridgeProcessor, BusFrameMessage +from pipecat.frames.frames import TextFrame +from pipecat.pipeline.pipeline import Pipeline +from pipecat.processors.frame_processor import FrameDirection +from pipecat.tests.utils import run_test + + +class TestBusBridgeProcessor(unittest.IsolatedAsyncioTestCase): + async def test_frames_sent_to_bus_not_passed_through(self): + """Non-lifecycle frames are sent to the bus, not passed through.""" + bus = AsyncQueueBus() + sent_to_bus = [] + original_send = bus.send + + async def capture_send(msg): + sent_to_bus.append(msg) + await original_send(msg) + + bus.send = capture_send + + processor = BusBridgeProcessor( + bus=bus, + task_name="test_task", + ) + pipeline = Pipeline([processor]) + + frames_to_send = [TextFrame(text="hello")] + + down, _ = await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=[], + ) + + # Frame NOT passed through downstream + text_frames = [f for f in down if isinstance(f, TextFrame)] + self.assertEqual(len(text_frames), 0) + + # Frame sent to bus + bus_frame_msgs = [m for m in sent_to_bus if isinstance(m, BusFrameMessage)] + self.assertEqual(len(bus_frame_msgs), 1) + self.assertEqual(bus_frame_msgs[0].frame.text, "hello") + self.assertEqual(bus_frame_msgs[0].source, "test_task") + self.assertEqual(bus_frame_msgs[0].direction, FrameDirection.DOWNSTREAM) + + async def test_lifecycle_frames_pass_through_not_sent_to_bus(self): + """Lifecycle frames pass through but are never sent to the bus.""" + bus = AsyncQueueBus() + sent_to_bus = [] + original_send = bus.send + + async def capture_send(msg): + sent_to_bus.append(msg) + await original_send(msg) + + bus.send = capture_send + + processor = BusBridgeProcessor( + bus=bus, + task_name="test_task", + ) + pipeline = Pipeline([processor]) + + # run_test sends StartFrame + frames_to_send + EndFrame + # TextFrame goes to bus (not downstream), lifecycle frames pass through + frames_to_send = [TextFrame(text="hello")] + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=[], + ) + + bus_frame_msgs = [m for m in sent_to_bus if isinstance(m, BusFrameMessage)] + # Only the TextFrame, not StartFrame or EndFrame + self.assertEqual(len(bus_frame_msgs), 1) + self.assertIsInstance(bus_frame_msgs[0].frame, TextFrame) + + async def test_exclude_frames_not_sent_to_bus(self): + """Excluded frame types pass through but are not sent to the bus.""" + bus = AsyncQueueBus() + sent_to_bus = [] + original_send = bus.send + + async def capture_send(msg): + sent_to_bus.append(msg) + await original_send(msg) + + bus.send = capture_send + + processor = BusBridgeProcessor( + bus=bus, + task_name="test_task", + exclude_frames=(TextFrame,), + ) + pipeline = Pipeline([processor]) + + frames_to_send = [TextFrame(text="excluded")] + expected_down_frames = [TextFrame] + + down, _ = await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + + # Frame passed through + self.assertEqual(len(down), 1) + self.assertEqual(down[0].text, "excluded") + + # But NOT sent to bus + bus_frame_msgs = [m for m in sent_to_bus if isinstance(m, BusFrameMessage)] + self.assertEqual(len(bus_frame_msgs), 0) + + async def test_bus_frame_injected_at_bridge(self): + """Frames from the bus are injected at the bridge position and + travel downstream alongside frames from later pipeline processors.""" + from pipecat.frames.frames import EndFrame + from pipecat.pipeline.runner import PipelineRunner + from pipecat.pipeline.task import PipelineTask + from pipecat.processors.frame_processor import FrameProcessor + + class AppendFrameProcessor(FrameProcessor): + """Appends a TextFrame for every TextFrame it sees.""" + + async def process_frame(self, frame, direction): + await super().process_frame(frame, direction) + await self.push_frame(frame, direction) + if isinstance(frame, TextFrame): + await self.push_frame(TextFrame(text="after_bridge"), direction) + + bus = AsyncQueueBus() + bridge = BusBridgeProcessor( + bus=bus, + task_name="main_task", + ) + pipeline = Pipeline([bridge, AppendFrameProcessor()]) + task = PipelineTask(pipeline, cancel_on_idle_timeout=False) + + received = [] + task.set_reached_downstream_filter((TextFrame,)) + + @task.event_handler("on_frame_reached_downstream") + async def on_frame(task, frame): + received.append(frame) + + msg = BusFrameMessage( + source="child_task", + frame=TextFrame(text="from_child"), + direction=FrameDirection.DOWNSTREAM, + ) + + async def inject_and_end(): + await asyncio.sleep(0.02) + await bridge.on_bus_message(msg) + await asyncio.sleep(0.02) + await task.queue_frame(EndFrame()) + + runner = PipelineRunner() + await asyncio.gather(runner.run(task), inject_and_end()) + + texts = [f.text for f in received if isinstance(f, TextFrame)] + self.assertIn("from_child", texts) + self.assertIn("after_bridge", texts) + + async def test_skips_own_frames(self): + """Bridge ignores bus frames from its own task.""" + bus = AsyncQueueBus() + processor = BusBridgeProcessor( + bus=bus, + task_name="test_task", + ) + + injected = [] + original_push = processor.push_frame + + async def capture_push(frame, direction=FrameDirection.DOWNSTREAM): + injected.append(frame) + await original_push(frame, direction) + + processor.push_frame = capture_push + + # Own frame should be ignored + msg = BusFrameMessage( + source="test_task", + frame=TextFrame(text="self"), + direction=FrameDirection.DOWNSTREAM, + ) + await processor.on_bus_message(msg) + + # Should not have injected anything + self.assertEqual(len(injected), 0) + + async def test_target_task_filtering(self): + """Bridge with target_task only accepts frames from that task.""" + bus = AsyncQueueBus() + processor = BusBridgeProcessor( + bus=bus, + task_name="main_task", + target_task="specific_child", + ) + + injected = [] + original_push = processor.push_frame + + async def capture_push(frame, direction=FrameDirection.DOWNSTREAM): + injected.append(frame) + await original_push(frame, direction) + + processor.push_frame = capture_push + + # Frame from wrong task — should be ignored + wrong_msg = BusFrameMessage( + source="other_child", + frame=TextFrame(text="wrong"), + direction=FrameDirection.DOWNSTREAM, + ) + await processor.on_bus_message(wrong_msg) + self.assertEqual(len(injected), 0) + + # Frame from correct task — should be injected + right_msg = BusFrameMessage( + source="specific_child", + frame=TextFrame(text="right"), + direction=FrameDirection.DOWNSTREAM, + ) + await processor.on_bus_message(right_msg) + self.assertEqual(len(injected), 1) + self.assertEqual(injected[0].text, "right") + + async def test_targeted_message_for_other_task_skipped(self): + """Bridge skips bus messages targeted at a different task.""" + bus = AsyncQueueBus() + processor = BusBridgeProcessor( + bus=bus, + task_name="main_task", + ) + + injected = [] + original_push = processor.push_frame + + async def capture_push(frame, direction=FrameDirection.DOWNSTREAM): + injected.append(frame) + await original_push(frame, direction) + + processor.push_frame = capture_push + + msg = BusFrameMessage( + source="child", + target="other_task", + frame=TextFrame(text="not_for_me"), + direction=FrameDirection.DOWNSTREAM, + ) + await processor.on_bus_message(msg) + self.assertEqual(len(injected), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bus.py b/tests/test_bus.py new file mode 100644 index 000000000..e6bac7be7 --- /dev/null +++ b/tests/test_bus.py @@ -0,0 +1,274 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import itertools +import unittest + +from pipecat.bus import ( + AsyncQueueBus, + BusCancelMessage, + BusCancelTaskMessage, + BusDataMessage, + BusJobCancelMessage, + BusSubscriber, +) +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + + +async def create_test_bus(): + """Create an AsyncQueueBus with a TaskManager for testing.""" + bus = AsyncQueueBus() + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + return bus, tm + + +class _CollectorSub(BusSubscriber): + """Subscriber that collects messages into a list.""" + + _counter = itertools.count() + + def __init__(self): + self._name = f"collector_{next(self._counter)}" + self.received = [] + + @property + def name(self) -> str: + return self._name + + async def on_bus_message(self, message): + self.received.append(message) + + +class TestAsyncQueueBus(unittest.IsolatedAsyncioTestCase): + async def test_send_delivers_to_subscriber(self): + """send() delivers a message to a subscriber.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + msg = BusDataMessage(source="task_a") + await bus.send(msg) + await asyncio.sleep(0.05) + await bus.stop() + + self.assertEqual(len(sub.received), 1) + self.assertIs(sub.received[0], msg) + + async def test_multiple_messages_in_order(self): + """Messages are dispatched in FIFO order.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + msgs = [BusDataMessage(source=f"task_{i}") for i in range(5)] + for m in msgs: + await bus.send(m) + await asyncio.sleep(0.1) + await bus.stop() + + self.assertEqual(len(sub.received), 5) + for sent, got in zip(msgs, sub.received): + self.assertIs(sent, got) + + async def test_start_stop_lifecycle(self): + """start() begins dispatch tasks, stop() cancels them cleanly.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + await bus.send(BusDataMessage(source="a")) + await asyncio.sleep(0.05) + self.assertEqual(len(sub.received), 1) + + await bus.stop() + + await bus.send(BusDataMessage(source="b")) + await asyncio.sleep(0.05) + # After stop, messages are not dispatched + self.assertEqual(len(sub.received), 1) + + +class TestBusSubscriber(unittest.IsolatedAsyncioTestCase): + async def test_subscribe_calls_on_bus_message(self): + """subscribe() delivers messages to subscriber's on_bus_message.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + msg = BusDataMessage(source="task_a") + await bus.send(msg) + await asyncio.sleep(0.05) + await bus.stop() + + self.assertEqual(len(sub.received), 1) + self.assertIs(sub.received[0], msg) + + async def test_multiple_subscribers_independent(self): + """Two subscribers each get every message on their own task.""" + bus, _ = await create_test_bus() + sub1 = _CollectorSub() + sub2 = _CollectorSub() + await bus.subscribe(sub1) + await bus.subscribe(sub2) + await bus.start() + + msg = BusDataMessage(source="task_a") + await bus.send(msg) + await asyncio.sleep(0.05) + await bus.stop() + + self.assertEqual(len(sub1.received), 1) + self.assertEqual(len(sub2.received), 1) + self.assertIs(sub1.received[0], msg) + self.assertIs(sub2.received[0], msg) + + async def test_unsubscribe_stops_delivery(self): + """unsubscribe() prevents further message delivery.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + await bus.send(BusDataMessage(source="a")) + await asyncio.sleep(0.05) + self.assertEqual(len(sub.received), 1) + + await bus.unsubscribe(sub) + await bus.send(BusDataMessage(source="b")) + await asyncio.sleep(0.05) + await bus.stop() + + self.assertEqual(len(sub.received), 1) + + async def test_slow_subscriber_does_not_block_others(self): + """A slow subscriber does not block a fast subscriber.""" + bus, _ = await create_test_bus() + fast_received = [] + fast_done = asyncio.Event() + + class SlowSub(BusSubscriber): + @property + def name(self) -> str: + return "slow" + + async def on_bus_message(self, message): + await asyncio.sleep(0.5) + + class FastSub(BusSubscriber): + @property + def name(self) -> str: + return "fast" + + async def on_bus_message(self, message): + fast_received.append(message) + fast_done.set() + + await bus.subscribe(SlowSub()) + await bus.subscribe(FastSub()) + await bus.start() + + await bus.send(BusDataMessage(source="a")) + await asyncio.wait_for(fast_done.wait(), timeout=0.1) + await bus.stop() + + self.assertEqual(len(fast_received), 1) + + +class TestBusMessagePriority(unittest.IsolatedAsyncioTestCase): + async def test_system_message_preempts_data_messages(self): + """System messages are delivered before data messages queued earlier.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + + # Queue messages before starting dispatch + for i in range(5): + await bus.send(BusDataMessage(source=f"data_{i}")) + await bus.send(BusCancelMessage(source="runner", reason="urgent")) + + await bus.start() + await asyncio.sleep(0.1) + await bus.stop() + + # System message should be first + self.assertIsInstance(sub.received[0], BusCancelMessage) + self.assertEqual(sub.received[0].source, "runner") + + async def test_data_messages_preserve_fifo_order(self): + """Data messages maintain FIFO order among themselves.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + for i in range(5): + await bus.send(BusDataMessage(source=f"task_{i}")) + await asyncio.sleep(0.1) + await bus.stop() + + sources = [m.source for m in sub.received] + self.assertEqual(sources, [f"task_{i}" for i in range(5)]) + + async def test_system_messages_preserve_fifo_order(self): + """Multiple system messages maintain FIFO order among themselves.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + await bus.start() + + await bus.send(BusCancelMessage(source="first", reason="a")) + await bus.send(BusCancelMessage(source="second", reason="b")) + await asyncio.sleep(0.1) + await bus.stop() + + self.assertEqual(sub.received[0].source, "first") + self.assertEqual(sub.received[1].source, "second") + + async def test_mixed_messages_system_first(self): + """When data and system messages are queued, all system come first.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + + # Queue before starting so all messages are pending + await bus.send(BusDataMessage(source="data_1")) + await bus.send(BusDataMessage(source="data_2")) + await bus.send(BusCancelMessage(source="cancel_1")) + await bus.send(BusDataMessage(source="data_3")) + await bus.send(BusCancelTaskMessage(source="cancel_2", target="task")) + + await bus.start() + await asyncio.sleep(0.1) + await bus.stop() + + sources = [m.source for m in sub.received] + self.assertEqual(sources, ["cancel_1", "cancel_2", "data_1", "data_2", "data_3"]) + + async def test_task_cancel_is_system_priority(self): + """BusJobCancelMessage has system priority.""" + bus, _ = await create_test_bus() + sub = _CollectorSub() + await bus.subscribe(sub) + + await bus.send(BusDataMessage(source="data")) + await bus.send(BusJobCancelMessage(source="parent", target="worker", job_id="t1")) + + await bus.start() + await asyncio.sleep(0.1) + await bus.stop() + + self.assertIsInstance(sub.received[0], BusJobCancelMessage) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bus_network.py b/tests/test_bus_network.py new file mode 100644 index 000000000..fdba47db4 --- /dev/null +++ b/tests/test_bus_network.py @@ -0,0 +1,35 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Tests for ``pipecat.bus.network`` lazy imports.""" + +import unittest + +import pipecat.bus.network as network + + +class TestNetworkPackageLazyImports(unittest.TestCase): + def test_pgmq_bus_lazy_import(self): + """``PgmqBus`` is resolved on first attribute access.""" + from pipecat.bus.network.pgmq import PgmqBus as Direct + + self.assertIs(network.PgmqBus, Direct) + + def test_redis_bus_lazy_import(self): + """``RedisBus`` is resolved on first attribute access.""" + from pipecat.bus.network.redis import RedisBus as Direct + + self.assertIs(network.RedisBus, Direct) + + def test_unknown_attribute_raises(self): + """Unknown attributes raise a clear ``AttributeError``.""" + with self.assertRaises(AttributeError) as ctx: + network.NotARealBus # noqa: B018 + self.assertIn("NotARealBus", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_frame_adapters.py b/tests/test_frame_adapters.py new file mode 100644 index 000000000..43627f528 --- /dev/null +++ b/tests/test_frame_adapters.py @@ -0,0 +1,163 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import unittest + +from pipecat.bus.messages import BusFrameMessage +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.frames.frames import LLMContextFrame, TextFrame, TranscriptionFrame +from pipecat.processors.aggregators.llm_context import ( + LLMContext, + LLMSpecificMessage, + NotGiven, +) +from pipecat.processors.frame_processor import FrameDirection + + +class TestTextFrameRoundTrip(unittest.TestCase): + def setUp(self): + self.serializer = JSONMessageSerializer() + + def test_round_trip(self): + msg = BusFrameMessage( + source="a", + frame=TextFrame(text="hello world"), + direction=FrameDirection.DOWNSTREAM, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored.frame, TextFrame) + self.assertEqual(restored.frame.text, "hello world") + + +class TestTranscriptionFrameRoundTrip(unittest.TestCase): + def setUp(self): + self.serializer = JSONMessageSerializer() + + def test_round_trip_basic(self): + msg = BusFrameMessage( + source="a", + frame=TranscriptionFrame( + text="hello", + user_id="user-1", + timestamp="2026-03-17T00:00:00Z", + ), + direction=FrameDirection.DOWNSTREAM, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored.frame, TranscriptionFrame) + self.assertEqual(restored.frame.text, "hello") + self.assertEqual(restored.frame.user_id, "user-1") + self.assertEqual(restored.frame.timestamp, "2026-03-17T00:00:00Z") + self.assertIsNone(restored.frame.language) + self.assertFalse(restored.frame.finalized) + + def test_round_trip_with_language(self): + from pipecat.transcriptions.language import Language + + msg = BusFrameMessage( + source="a", + frame=TranscriptionFrame( + text="hola", + user_id="user-1", + timestamp="2026-03-17T00:00:00Z", + language=Language.ES, + finalized=True, + ), + direction=FrameDirection.DOWNSTREAM, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertEqual(restored.frame.language, Language.ES) + self.assertTrue(restored.frame.finalized) + + +class TestLLMContextFrameRoundTrip(unittest.TestCase): + def setUp(self): + self.serializer = JSONMessageSerializer() + + def _round_trip_context(self, ctx): + msg = BusFrameMessage( + source="a", + frame=LLMContextFrame(context=ctx), + direction=FrameDirection.DOWNSTREAM, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + self.assertIsInstance(restored.frame, LLMContextFrame) + return restored.frame.context + + def test_round_trip_messages_only(self): + ctx = LLMContext( + messages=[ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "hello"}, + ] + ) + restored = self._round_trip_context(ctx) + + self.assertEqual(len(restored.messages), 2) + self.assertEqual(restored.messages[0]["role"], "system") + self.assertEqual(restored.messages[1]["content"], "hello") + self.assertIsInstance(restored.tools, NotGiven) + self.assertIsInstance(restored.tool_choice, NotGiven) + + def test_round_trip_with_specific_message(self): + ctx = LLMContext( + messages=[ + {"role": "user", "content": "hi"}, + LLMSpecificMessage(llm="anthropic", message={"custom": "data"}), + ] + ) + restored = self._round_trip_context(ctx) + + self.assertEqual(len(restored.messages), 2) + self.assertIsInstance(restored.messages[0], dict) + self.assertIsInstance(restored.messages[1], LLMSpecificMessage) + self.assertEqual(restored.messages[1].llm, "anthropic") + self.assertEqual(restored.messages[1].message, {"custom": "data"}) + + def test_round_trip_with_tools(self): + from pipecat.adapters.schemas.function_schema import FunctionSchema + from pipecat.adapters.schemas.tools_schema import ToolsSchema + + tools = ToolsSchema( + standard_tools=[ + FunctionSchema( + name="get_weather", + description="Get the weather", + properties={"location": {"type": "string"}}, + required=["location"], + ), + ] + ) + ctx = LLMContext( + messages=[{"role": "user", "content": "weather?"}], + tools=tools, + ) + restored = self._round_trip_context(ctx) + + self.assertNotIsInstance(restored.tools, NotGiven) + self.assertEqual(len(restored.tools.standard_tools), 1) + self.assertEqual(restored.tools.standard_tools[0].name, "get_weather") + self.assertEqual(restored.tools.standard_tools[0].required, ["location"]) + + def test_round_trip_with_tool_choice(self): + ctx = LLMContext( + messages=[{"role": "user", "content": "hi"}], + tool_choice="auto", + ) + restored = self._round_trip_context(ctx) + + self.assertEqual(restored.tool_choice, "auto") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_job_group.py b/tests/test_job_group.py new file mode 100644 index 000000000..210bb63d5 --- /dev/null +++ b/tests/test_job_group.py @@ -0,0 +1,739 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest + +from pipecat.bus import ( + AsyncQueueBus, + BusJobCancelMessage, + BusJobRequestMessage, + BusJobResponseMessage, +) +from pipecat.pipeline.base_task import BaseTask +from pipecat.pipeline.job_context import ( + JobError, + JobEvent, + JobGroupError, + JobGroupEvent, + JobStatus, +) +from pipecat.registry import TaskRegistry +from pipecat.registry.types import TaskReadyData +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + + +class StubTask(BaseTask): + pass + + +class JobWorkerTask(BaseTask): + """Worker that automatically responds to job requests via the bus.""" + + def __init__(self, name, *, response=None, status=JobStatus.COMPLETED): + super().__init__(name) + self._auto_response = response + self._auto_status = status + + async def on_job_request(self, message): + await super().on_job_request(message) + await self.send_job_response(message.job_id, self._auto_response, status=self._auto_status) + + +class UrgentJobWorkerTask(BaseTask): + """Worker that responds urgently to job requests.""" + + def __init__(self, name, *, response=None, status=JobStatus.COMPLETED): + super().__init__(name) + self._auto_response = response + self._auto_status = status + + async def on_job_request(self, message): + await super().on_job_request(message) + await self.send_job_response( + message.job_id, self._auto_response, status=self._auto_status, urgent=True + ) + + +class UpdatingWorkerTask(BaseTask): + """Worker that sends updates before responding.""" + + def __init__(self, name, *, updates, response=None): + super().__init__(name) + self._updates = updates + self._auto_response = response + + async def on_job_request(self, message): + await super().on_job_request(message) + for update in self._updates: + await self.send_job_update(message.job_id, update) + await self.send_job_response(message.job_id, self._auto_response) + + +class StreamingWorkerTask(BaseTask): + """Worker that streams data before responding.""" + + def __init__(self, name, *, chunks, response=None): + super().__init__(name) + self._chunks = chunks + self._auto_response = response + + async def on_job_request(self, message): + await super().on_job_request(message) + await self.send_job_stream_start(message.job_id, {"content_type": "text"}) + for chunk in self._chunks: + await self.send_job_stream_data(message.job_id, chunk) + await self.send_job_stream_end(message.job_id, self._auto_response) + + +class SlowWorkerTask(BaseTask): + """Worker that blocks during job execution until cancelled.""" + + def __init__(self, name): + super().__init__(name) + self.started = asyncio.Event() + self.was_cancelled = False + + async def on_job_request(self, message): + await super().on_job_request(message) + self.started.set() + try: + # Block until cancelled + await asyncio.Event().wait() + except asyncio.CancelledError: + self.was_cancelled = True + + +async def create_test_env(): + bus = AsyncQueueBus() + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + registry = TaskRegistry(runner_name="test-runner") + return bus, tm, registry + + +async def setup_task(bus, registry, task): + """Subscribe a task to the bus and register it as ready.""" + task.attach(registry=registry, bus=bus) + await task.setup(bus.task_manager) + await bus.subscribe(task) + await registry.register(TaskReadyData(task_name=task.name, runner="test-runner")) + + +def capture_bus(bus): + sent = [] + original_send = bus.send + + async def capture_send(message): + sent.append(message) + await original_send(message) + + bus.send = capture_send + return sent + + +class TestJobGroupContext(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm, self.registry = await create_test_env() + + async def asyncTearDown(self): + await self.bus.stop() + + async def test_job_group_collects_responses(self): + """job_group() context manager collects all responses.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + w1 = JobWorkerTask("w1", response={"a": 1}) + w2 = JobWorkerTask("w2", response={"b": 2}) + await setup_task(self.bus, self.registry, w1) + await setup_task(self.bus, self.registry, w2) + + async with parent.job_group("w1", "w2", payload={"work": True}) as tg: + pass + + self.assertEqual(tg.responses, {"w1": {"a": 1}, "w2": {"b": 2}}) + + async def test_job_group_sends_request(self): + """job_group() sends BusJobRequestMessage to each task.""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + w1 = JobWorkerTask("w1", response={"ok": True}) + await setup_task(self.bus, self.registry, w1) + + async with parent.job_group("w1", payload={"data": 1}) as tg: + pass + + request_msgs = [m for m in sent if isinstance(m, BusJobRequestMessage)] + self.assertEqual(len(request_msgs), 1) + self.assertEqual(request_msgs[0].payload, {"data": 1}) + + async def test_job_group_raises_on_cancel(self): + """job_group() raises JobGroupError when group is cancelled.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + # StubTask doesn't auto-respond, so we can cancel manually + worker = StubTask("worker") + await setup_task(self.bus, self.registry, worker) + + with self.assertRaises(JobGroupError) as ctx: + async with parent.job_group("worker") as tg: + asyncio.ensure_future(parent.cancel_job_group(tg.job_id, reason="manual cancel")) + + # Let the event loop schedule handler tasks spawned by the bus + await asyncio.sleep(0) + self.assertIn("manual cancel", str(ctx.exception)) + + async def test_job_group_raises_on_timeout(self): + """job_group() raises JobGroupError on timeout.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + # StubTask doesn't auto-respond, so timeout fires + worker = StubTask("worker") + await setup_task(self.bus, self.registry, worker) + + with self.assertRaises(JobGroupError) as ctx: + async with parent.job_group("worker", timeout=0.05) as tg: + pass + + self.assertIn("timeout", str(ctx.exception)) + + async def test_job_group_raises_on_ready_timeout(self): + """job_group() raises JobGroupError when tasks aren't ready in time.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + # "ghost" is never registered, so the ready-wait times out + with self.assertRaises(JobGroupError) as ctx: + async with parent.job_group("ghost", timeout=0.05) as tg: + pass + + self.assertIn("not ready", str(ctx.exception)) + + async def test_job_group_cancels_on_block_exception(self): + """job_group() cancels remaining jobs when the block raises.""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = StubTask("worker") + await setup_task(self.bus, self.registry, worker) + + with self.assertRaises(ValueError): + async with parent.job_group("worker") as tg: + raise ValueError("something went wrong") + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].reason, "context exited with error") + + async def test_job_group_raises_on_worker_error(self): + """job_group() raises JobGroupError when a worker errors with cancel_on_error.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={"error": "failed"}, status=JobStatus.ERROR) + await setup_task(self.bus, self.registry, worker) + + with self.assertRaises(JobGroupError): + async with parent.job_group("worker") as tg: + pass + + # Error response is tracked in partial responses + self.assertEqual(tg.responses, {"worker": {"error": "failed"}}) + + async def test_job_group_on_job_error_fires(self): + """on_job_error fires when a worker errors with cancel_on_error.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={"error": "boom"}, status=JobStatus.ERROR) + await setup_task(self.bus, self.registry, worker) + + errors = [] + + @parent.event_handler("on_job_error") + async def on_error(task, message): + errors.append(message) + + with self.assertRaises(JobGroupError): + async with parent.job_group("worker") as tg: + pass + + self.assertEqual(len(errors), 1) + self.assertEqual(errors[0].source, "worker") + self.assertEqual(errors[0].response, {"error": "boom"}) + self.assertEqual(errors[0].status, JobStatus.ERROR) + + async def test_job_group_partial_responses_on_error(self): + """Partial responses from successful workers are available after error.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + # w1 responds successfully, w2 responds with error. + # Order depends on bus dispatch, but both are registered. + w1 = JobWorkerTask("w1", response={"ok": True}) + w2 = JobWorkerTask("w2", response={"error": "fail"}, status=JobStatus.ERROR) + await setup_task(self.bus, self.registry, w1) + await setup_task(self.bus, self.registry, w2) + + with self.assertRaises(JobGroupError): + async with parent.job_group("w1", "w2") as tg: + pass + + # w2's error response should be in partial responses + self.assertIn("w2", tg.responses) + self.assertEqual(tg.responses["w2"], {"error": "fail"}) + + async def test_job_group_job_id_available(self): + """job_id is available inside the async with block.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={}) + await setup_task(self.bus, self.registry, worker) + + captured_job_id = None + + async with parent.job_group("worker") as tg: + captured_job_id = tg.job_id + + self.assertIsNotNone(captured_job_id) + self.assertEqual(captured_job_id, tg.job_id) + + async def test_job_group_on_job_completed_still_fires(self): + """on_job_completed callback still fires when using job_group().""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("w1", response={"ok": True}) + await setup_task(self.bus, self.registry, worker) + + completed = [] + + @parent.event_handler("on_job_completed") + async def on_completed(task, result): + completed.append(result) + + async with parent.job_group("w1") as tg: + pass + + self.assertEqual(len(completed), 1) + self.assertEqual(completed[0].responses, {"w1": {"ok": True}}) + + async def test_job_group_iterates_updates(self): + """async for yields update events from workers.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = UpdatingWorkerTask( + "worker", + updates=[{"progress": 25}, {"progress": 75}], + response={"result": "done"}, + ) + await setup_task(self.bus, self.registry, worker) + + events = [] + async with parent.job_group("worker", payload={"work": True}) as tg: + async for event in tg: + events.append(event) + + self.assertEqual(len(events), 2) + self.assertEqual(events[0].type, JobGroupEvent.UPDATE) + self.assertEqual(events[0].task_name, "worker") + self.assertEqual(events[0].data, {"progress": 25}) + self.assertEqual(events[1].data, {"progress": 75}) + self.assertEqual(tg.responses, {"worker": {"result": "done"}}) + + async def test_job_group_iterates_stream_events(self): + """async for yields stream_start, stream_data, and stream_end events.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = StreamingWorkerTask( + "worker", + chunks=[{"text": "hello "}, {"text": "world"}], + response={"final": True}, + ) + await setup_task(self.bus, self.registry, worker) + + events = [] + async with parent.job_group("worker") as tg: + async for event in tg: + events.append(event) + + types = [e.type for e in events] + self.assertEqual( + types, + [ + JobGroupEvent.STREAM_START, + JobGroupEvent.STREAM_DATA, + JobGroupEvent.STREAM_DATA, + JobGroupEvent.STREAM_END, + ], + ) + self.assertEqual(events[0].data, {"content_type": "text"}) + self.assertEqual(events[1].data, {"text": "hello "}) + self.assertEqual(events[2].data, {"text": "world"}) + self.assertEqual(events[3].data, {"final": True}) + + async def test_job_group_iterates_mixed_events(self): + """async for yields events from multiple workers interleaved.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + w1 = UpdatingWorkerTask("w1", updates=[{"status": "working"}], response={"a": 1}) + w2 = JobWorkerTask("w2", response={"b": 2}) + await setup_task(self.bus, self.registry, w1) + await setup_task(self.bus, self.registry, w2) + + events = [] + async with parent.job_group("w1", "w2") as tg: + async for event in tg: + events.append(event) + + self.assertEqual(len(events), 1) + self.assertEqual(events[0].type, JobGroupEvent.UPDATE) + self.assertEqual(events[0].task_name, "w1") + self.assertEqual(tg.responses, {"w1": {"a": 1}, "w2": {"b": 2}}) + + async def test_job_group_no_iteration_still_works(self): + """job_group() works without iterating (pass body).""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = UpdatingWorkerTask("worker", updates=[{"progress": 50}], response={"ok": True}) + await setup_task(self.bus, self.registry, worker) + + async with parent.job_group("worker") as tg: + pass + + self.assertEqual(tg.responses, {"worker": {"ok": True}}) + + async def test_urgent_job_response_collected(self): + """Urgent job responses are collected like normal responses.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = UrgentJobWorkerTask("worker", response={"urgent": True}) + await setup_task(self.bus, self.registry, worker) + + async with parent.job_group("worker") as tg: + pass + + self.assertEqual(tg.responses, {"worker": {"urgent": True}}) + + async def test_urgent_job_response_triggers_on_job_error(self): + """Urgent error response triggers on_job_error and cancels group.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = UrgentJobWorkerTask( + "worker", response={"error": "critical"}, status=JobStatus.ERROR + ) + await setup_task(self.bus, self.registry, worker) + + errors = [] + + @parent.event_handler("on_job_error") + async def on_error(task, message): + errors.append(message) + + with self.assertRaises(JobGroupError): + async with parent.job_group("worker") as tg: + pass + + self.assertEqual(len(errors), 1) + self.assertEqual(errors[0].response, {"error": "critical"}) + + async def test_urgent_response_has_system_priority(self): + """Urgent job response is delivered before normal data messages.""" + from pipecat.bus import BusDataMessage, BusJobResponseUrgentMessage + + bus = self.bus + + # Queue data messages first, then an urgent response + parent = StubTask("parent") + await setup_task(bus, self.registry, parent) + + received = [] + + @parent.event_handler("on_bus_message") + async def on_msg(task, message): + received.append(message) + + # Send data messages before starting dispatch + for i in range(3): + await bus.send(BusDataMessage(source=f"data_{i}")) + await bus.send( + BusJobResponseUrgentMessage( + source="worker", target="parent", job_id="t1", status=JobStatus.COMPLETED + ) + ) + + # Start dispatch — urgent should arrive first + # bus is already started in setUp; restart by stopping and starting + await bus.stop() + await bus.start() + await asyncio.sleep(0.1) + await bus.stop() + + # First message should be the urgent response + self.assertIsInstance(received[0], BusJobResponseUrgentMessage) + + +class TestJobContext(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm, self.registry = await create_test_env() + + async def asyncTearDown(self): + await self.bus.stop() + + async def test_job_collects_response(self): + """job() context manager collects the worker's response.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={"result": 42}) + await setup_task(self.bus, self.registry, worker) + + async with parent.job("worker", payload={"x": 1}) as t: + pass + + self.assertEqual(t.response, {"result": 42}) + + async def test_job_sends_request(self): + """job() sends a BusJobRequestMessage to the task.""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={"ok": True}) + await setup_task(self.bus, self.registry, worker) + + async with parent.job("worker") as t: + pass + + request_msgs = [m for m in sent if isinstance(m, BusJobRequestMessage)] + self.assertEqual(len(request_msgs), 1) + self.assertEqual(request_msgs[0].target, "worker") + self.assertEqual(request_msgs[0].source, "parent") + + async def test_job_iterates_events(self): + """job() yields intermediate events via async for.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = UpdatingWorkerTask("worker", updates=[{"progress": 50}], response={"done": True}) + await setup_task(self.bus, self.registry, worker) + + events = [] + async with parent.job("worker") as t: + async for event in t: + events.append(event) + + self.assertEqual(len(events), 1) + self.assertEqual(events[0].type, JobEvent.UPDATE) + self.assertIsInstance(events[0], JobEvent) + self.assertEqual(t.response, {"done": True}) + + async def test_job_streams_data(self): + """job() yields stream events from a streaming worker.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = StreamingWorkerTask( + "worker", + chunks=[{"text": "hello"}, {"text": "world"}], + response={"ok": True}, + ) + await setup_task(self.bus, self.registry, worker) + + events = [] + async with parent.job("worker") as t: + async for event in t: + events.append(event) + + types = [e.type for e in events] + self.assertEqual( + types, + [ + JobEvent.STREAM_START, + JobEvent.STREAM_DATA, + JobEvent.STREAM_DATA, + JobEvent.STREAM_END, + ], + ) + self.assertEqual(t.response, {"ok": True}) + + async def test_job_cancels_on_exception(self): + """job() cancels the job if the block raises.""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = StubTask("worker") + await setup_task(self.bus, self.registry, worker) + + with self.assertRaises(ValueError): + async with parent.job("worker") as t: + raise ValueError("something went wrong") + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].reason, "context exited with error") + + async def test_job_raises_on_worker_error(self): + """job() raises JobError when the worker errors.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={"error": "boom"}, status=JobStatus.ERROR) + await setup_task(self.bus, self.registry, worker) + + with self.assertRaises(JobError): + async with parent.job("worker") as t: + pass + + async def test_job_exposes_job_id(self): + """job() exposes the job_id inside the context.""" + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = JobWorkerTask("worker", response={"ok": True}) + await setup_task(self.bus, self.registry, worker) + + async with parent.job("worker") as t: + self.assertIsInstance(t.job_id, str) + self.assertTrue(len(t.job_id) > 0) + + async def test_job_group_cancels_on_cancelled_error(self): + """job_group() cancels workers when CancelledError is raised (e.g. tool interruption).""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = StubTask("worker") + await setup_task(self.bus, self.registry, worker) + + async def run_job_group(): + async with parent.job_group("worker") as tg: + # Simulate tool cancellation while waiting + raise asyncio.CancelledError() + + task = asyncio.create_task(run_job_group()) + with self.assertRaises(asyncio.CancelledError): + await task + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].reason, "context exited with error") + # Job group should be cleaned up + self.assertEqual(len(parent.job_groups), 0) + + async def test_job_cancels_on_cancelled_error(self): + """job() cancels the worker when CancelledError is raised (e.g. tool interruption).""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = StubTask("worker") + await setup_task(self.bus, self.registry, worker) + + async def run_job(): + async with parent.job("worker") as t: + raise asyncio.CancelledError() + + task = asyncio.create_task(run_job()) + with self.assertRaises(asyncio.CancelledError): + await task + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].reason, "context exited with error") + self.assertEqual(len(parent.job_groups), 0) + + async def test_fire_and_forget_jobs_cancelled_manually(self): + """User cancels fire-and-forget jobs in a CancelledError handler.""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + w1 = StubTask("w1") + w2 = StubTask("w2") + await setup_task(self.bus, self.registry, w1) + await setup_task(self.bus, self.registry, w2) + + # Simulate what a user would do in on_function_calls_cancelled: + # track job IDs and cancel only the ones started here + try: + job_ids = [] + job_ids.append(await parent.request_job("w1", payload={"job": 1})) + job_ids.append(await parent.request_job("w2", payload={"job": 2})) + self.assertEqual(len(parent.job_groups), 2) + raise asyncio.CancelledError() + except asyncio.CancelledError: + for jid in job_ids: + await parent.cancel_job_group(jid, reason="tool cancelled") + + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 2) + cancelled_targets = {m.target for m in cancel_msgs} + self.assertEqual(cancelled_targets, {"w1", "w2"}) + for m in cancel_msgs: + self.assertEqual(m.reason, "tool cancelled") + self.assertEqual(len(parent.job_groups), 0) + + async def test_cancel_interrupts_running_handler(self): + """Cancelling a job interrupts a handler that is currently executing.""" + sent = capture_bus(self.bus) + + parent = StubTask("parent") + await setup_task(self.bus, self.registry, parent) + + worker = SlowWorkerTask("worker") + await setup_task(self.bus, self.registry, worker) + + job_id = await parent.request_job("worker", payload={"job": 1}) + + # Wait for the worker to start executing + await asyncio.wait_for(worker.started.wait(), timeout=2.0) + + # Cancel while the handler is blocked + await parent.cancel_job_group(job_id, reason="no longer needed") + + # Give the event loop time to process the cancellation + await asyncio.sleep(0.1) + + # The handler should have been cancelled + self.assertTrue(worker.was_cancelled) + + # A cancel message should have been sent + cancel_msgs = [m for m in sent if isinstance(m, BusJobCancelMessage)] + self.assertEqual(len(cancel_msgs), 1) + self.assertEqual(cancel_msgs[0].reason, "no longer needed") + + # A CANCELLED response should have been sent back + response_msgs = [m for m in sent if isinstance(m, BusJobResponseMessage)] + self.assertEqual(len(response_msgs), 1) + self.assertEqual(response_msgs[0].status, JobStatus.CANCELLED) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_llm_task.py b/tests/test_llm_task.py new file mode 100644 index 000000000..6e59cefb1 --- /dev/null +++ b/tests/test_llm_task.py @@ -0,0 +1,285 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest +from unittest.mock import MagicMock + +from pipecat.frames.frames import LLMMessagesAppendFrame +from pipecat.pipeline.task import PipelineTask +from pipecat.processors.frame_processor import FrameDirection +from pipecat.tasks.llm import LLMTask, tool +from pipecat.tasks.llm.llm_task import PipelineFlushFrame + + +def _create_task(): + """Create a StubLLMTask with mocked parent queue_frame for testing.""" + + class StubLLMTask(LLMTask): + @tool + async def fast_tool(self, params): + """A quick tool.""" + await params.result_callback("done") + + @tool + async def slow_tool(self, params, delay: float): + """A tool that blocks on an event for coordination.""" + await params.result_callback("done") + + llm = MagicMock() + llm.register_direct_function = MagicMock() + task = StubLLMTask("test_task", llm=llm, bridged=()) + + # Capture frames passed to PipelineTask.queue_frame (i.e. super().queue_frame). + # Auto-set _flush_done when PipelineFlushFrame is queued, simulating the flush + # round-trip through the pipeline. + delivered: list[tuple] = [] + original_pt_queue_frame = PipelineTask.queue_frame + + async def class_replacement(self, frame, direction=FrameDirection.DOWNSTREAM): + # Only intercept for this specific instance; otherwise fall through. + if self is task: + delivered.append((frame, direction)) + if isinstance(frame, PipelineFlushFrame): + task._flush_done.set() + return + await original_pt_queue_frame(self, frame, direction) + + PipelineTask.queue_frame = class_replacement + task._restore_pt_queue_frame = lambda: setattr( + PipelineTask, "queue_frame", original_pt_queue_frame + ) + + task._delivered_frames = delivered + return task + + +def _get_delivered_frames(task): + """Extract non-flush frames delivered to the underlying pipeline.""" + return [ + (frame, direction) + for frame, direction in task._delivered_frames + if not isinstance(frame, PipelineFlushFrame) + ] + + +def _make_frame(content: str, run_llm: bool = True) -> LLMMessagesAppendFrame: + return LLMMessagesAppendFrame(messages=[{"role": "user", "content": content}], run_llm=run_llm) + + +class TestToolCallTracking(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self._tasks = [] + + def _track(self, task): + self._tasks.append(task) + return task + + def tearDown(self): + for task in self._tasks: + restore = getattr(task, "_restore_pt_queue_frame", None) + if restore: + restore() + + async def test_tool_call_active_initially_false(self): + task = self._track(_create_task()) + self.assertFalse(task.tool_call_active) + + async def test_tool_call_active_during_execution(self): + """tool_call_active is True while a tool is running.""" + task = self._track(_create_task()) + observed = [] + + @tool + async def gated_tool(self, params): + """Waits on gate.""" + observed.append(task.tool_call_active) + + wrapped = task._track_tool_call(gated_tool.__get__(task)) + params = MagicMock() + await wrapped(params) + + self.assertTrue(observed[0]) + self.assertFalse(task.tool_call_active) + + async def test_queue_frame_delivers_immediately_when_idle(self): + """queue_frame delivers immediately when no tools are in-flight.""" + task = self._track(_create_task()) + frame = _make_frame("hello") + + await task.queue_frame(frame) + + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 1) + self.assertIs(delivered[0][0], frame) + + async def test_queue_frame_defers_when_tool_active(self): + """queue_frame defers delivery when a tool is in-flight.""" + task = self._track(_create_task()) + task._tool_call_inflight = 1 + + frame = _make_frame("deferred") + await task.queue_frame(frame) + + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 0) + self.assertEqual(len(task._deferred_frames), 1) + self.assertEqual(task._deferred_frames[0], (frame, FrameDirection.DOWNSTREAM)) + + async def test_deferred_frames_flush_when_tool_completes(self): + """Frames deferred during a tool call are delivered when it finishes.""" + task = self._track(_create_task()) + gate = asyncio.Event() + frame = _make_frame("event data") + + @tool + async def blocking_tool(self, params): + """Blocks until gate is set.""" + await gate.wait() + + wrapped = task._track_tool_call(blocking_tool.__get__(task)) + params = MagicMock() + + runner_task = asyncio.create_task(wrapped(params)) + await asyncio.sleep(0) + + await task.queue_frame(frame) + self.assertEqual(len(_get_delivered_frames(task)), 0) + + gate.set() + await runner_task + + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 1) + self.assertIs(delivered[0][0], frame) + + async def test_concurrent_tools_flush_only_when_all_done(self): + """With two parallel tools, flush happens only when the last one completes.""" + task = self._track(_create_task()) + gate_a = asyncio.Event() + gate_b = asyncio.Event() + + @tool + async def tool_a(self, params): + """First tool.""" + await gate_a.wait() + + @tool + async def tool_b(self, params): + """Second tool.""" + await gate_b.wait() + + wrapped_a = task._track_tool_call(tool_a.__get__(task)) + wrapped_b = task._track_tool_call(tool_b.__get__(task)) + params = MagicMock() + + task_a = asyncio.create_task(wrapped_a(params)) + task_b = asyncio.create_task(wrapped_b(params)) + await asyncio.sleep(0) + + self.assertEqual(task._tool_call_inflight, 2) + + frame = _make_frame("queued") + await task.queue_frame(frame) + + # First tool finishes — frame still deferred (second tool running) + gate_a.set() + await task_a + self.assertEqual(task._tool_call_inflight, 1) + self.assertEqual(len(_get_delivered_frames(task)), 0) + + # Second tool finishes — NOW flush + gate_b.set() + await task_b + self.assertEqual(task._tool_call_inflight, 0) + + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 1) + self.assertIs(delivered[0][0], frame) + + async def test_queue_frame_preserves_frame_attributes(self): + """Frame attributes like run_llm are preserved through defer and flush.""" + task = self._track(_create_task()) + gate = asyncio.Event() + + @tool + async def blocking_tool(self, params): + """Blocks.""" + await gate.wait() + + wrapped = task._track_tool_call(blocking_tool.__get__(task)) + params = MagicMock() + + runner_task = asyncio.create_task(wrapped(params)) + await asyncio.sleep(0) + + frame = _make_frame("no inference", run_llm=False) + await task.queue_frame(frame) + + gate.set() + await runner_task + + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 1) + self.assertFalse(delivered[0][0].run_llm) + + async def test_multiple_deferred_frames_flush_in_order(self): + """Multiple deferred frames are delivered in FIFO order.""" + task = self._track(_create_task()) + gate = asyncio.Event() + + @tool + async def blocking_tool(self, params): + """Blocks.""" + await gate.wait() + + wrapped = task._track_tool_call(blocking_tool.__get__(task)) + params = MagicMock() + + runner_task = asyncio.create_task(wrapped(params)) + await asyncio.sleep(0) + + frame_a = _make_frame("first", run_llm=False) + frame_b = _make_frame("second", run_llm=True) + await task.queue_frame(frame_a) + await task.queue_frame(frame_b) + + gate.set() + await runner_task + + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 2) + self.assertIs(delivered[0][0], frame_a) + self.assertIs(delivered[1][0], frame_b) + + async def test_tool_error_still_decrements_and_flushes(self): + """If a tool raises, the counter still decrements and deferred frames flush.""" + task = self._track(_create_task()) + + @tool + async def failing_tool(self, params): + """Always fails.""" + raise ValueError("boom") + + wrapped = task._track_tool_call(failing_tool.__get__(task)) + params = MagicMock() + + frame = _make_frame("recover") + task._tool_call_inflight = 1 + await task.queue_frame(frame) + task._tool_call_inflight = 0 + + with self.assertRaises(ValueError): + await wrapped(params) + + self.assertFalse(task.tool_call_active) + delivered = _get_delivered_frames(task) + self.assertEqual(len(delivered), 1) + self.assertIs(delivered[0][0], frame) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_messages.py b/tests/test_messages.py new file mode 100644 index 000000000..7bf75f672 --- /dev/null +++ b/tests/test_messages.py @@ -0,0 +1,97 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest + +from pipecat.bus import ( + AsyncQueueBus, + BusDataMessage, + BusFrameMessage, + BusSubscriber, +) +from pipecat.frames.frames import TextFrame +from pipecat.processors.frame_processor import FrameDirection +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + + +class TestBusMessageRouting(unittest.IsolatedAsyncioTestCase): + async def test_broadcast_message_no_target(self): + """A BusMessage with no target is broadcast (target is None).""" + msg = BusDataMessage(source="task_a") + self.assertIsNone(msg.target) + + async def test_targeted_message(self): + """A BusMessage with target set is only for that task.""" + msg = BusDataMessage(source="task_a", target="task_b") + self.assertEqual(msg.target, "task_b") + + async def test_broadcast_reaches_all_subscribers(self): + """Broadcast messages (no target) reach all subscribers.""" + bus = AsyncQueueBus() + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + received_a = [] + received_b = [] + + class SubA(BusSubscriber): + @property + def name(self) -> str: + return "sub_a" + + async def on_bus_message(self, message): + received_a.append(message) + + class SubB(BusSubscriber): + @property + def name(self) -> str: + return "sub_b" + + async def on_bus_message(self, message): + received_b.append(message) + + await bus.subscribe(SubA()) + await bus.subscribe(SubB()) + + await bus.start() + msg = BusDataMessage(source="task_x") # no target + await bus.send(msg) + await asyncio.sleep(0.05) + await bus.stop() + + # Both subscribers see the broadcast + self.assertEqual(len(received_a), 1) + self.assertEqual(len(received_b), 1) + + async def test_bus_frame_message_wraps_frame(self): + """BusFrameMessage wraps a frame with source and direction.""" + frame = TextFrame(text="hello") + msg = BusFrameMessage( + source="task_a", + frame=frame, + direction=FrameDirection.DOWNSTREAM, + ) + self.assertIs(msg.frame, frame) + self.assertEqual(msg.direction, FrameDirection.DOWNSTREAM) + self.assertEqual(msg.source, "task_a") + self.assertIsNone(msg.target) + + async def test_bus_frame_message_with_target(self): + """BusFrameMessage can carry a target for directed delivery.""" + frame = TextFrame(text="hello") + msg = BusFrameMessage( + source="task_a", + target="task_b", + frame=frame, + direction=FrameDirection.DOWNSTREAM, + ) + self.assertEqual(msg.target, "task_b") + self.assertIs(msg.frame, frame) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pgmq_backends.py b/tests/test_pgmq_backends.py new file mode 100644 index 000000000..ad2732a5e --- /dev/null +++ b/tests/test_pgmq_backends.py @@ -0,0 +1,266 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Unit tests for the PgmqBus backend layer. + +Covers: + +- ``IsolatedPgmqBackend`` issues the expected ``public.bus_*`` SQL with the + expected positional args, and normalizes the read result shape. +- ``PgmqBus`` correctly delegates to an injected backend (the + ``backend=`` constructor path) and rejects ambiguous / empty construction. + +``DirectPgmqBackend`` is exercised end-to-end via ``test_pgmq_bus.py`` +against an in-memory ``FakePgmq``; no separate tests for it here. +""" + +import asyncio +import json +import unittest +from dataclasses import dataclass + +from pipecat.bus import BusDataMessage +from pipecat.bus.network.pgmq import PgmqBus +from pipecat.bus.network.pgmq_backends import ( + BackendMessage, + IsolatedPgmqBackend, + PgmqBackend, +) +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + +# --------------------------------------------------------------------------- +# IsolatedPgmqBackend: assert it speaks the right SQL to asyncpg. +# --------------------------------------------------------------------------- + + +@dataclass +class _Call: + method: str # 'fetchval' | 'execute' | 'fetch' + sql: str + args: tuple + + +class _FakeConn: + def __init__(self, recorder, fetchval_return=None, fetch_return=None): + self._recorder = recorder + self._fetchval_return = fetchval_return + self._fetch_return = fetch_return or [] + + async def fetchval(self, sql, *args): + self._recorder.append(_Call("fetchval", sql, args)) + return self._fetchval_return + + async def execute(self, sql, *args): + self._recorder.append(_Call("execute", sql, args)) + return "OK" + + async def fetch(self, sql, *args): + self._recorder.append(_Call("fetch", sql, args)) + return self._fetch_return + + +class _FakePool: + def __init__(self, fetchval_return=None, fetch_return=None): + self.calls: list[_Call] = [] + self._fetchval_return = fetchval_return + self._fetch_return = fetch_return or [] + + def acquire(self): + recorder = self.calls + fetchval_return = self._fetchval_return + fetch_return = self._fetch_return + + class _Ctx: + async def __aenter__(self_inner): + return _FakeConn(recorder, fetchval_return, fetch_return) + + async def __aexit__(self_inner, exc_type, exc, tb): + return False + + return _Ctx() + + +class TestIsolatedPgmqBackend(unittest.IsolatedAsyncioTestCase): + async def test_join_calls_bus_join_and_returns_queue_name(self): + pool = _FakePool(fetchval_return="q_abc123") + backend = IsolatedPgmqBackend(pool=pool) + + queue = await backend.join("ch_xyz") + + self.assertEqual(queue, "q_abc123") + self.assertEqual(len(pool.calls), 1) + self.assertIn("public.bus_join", pool.calls[0].sql) + self.assertEqual(pool.calls[0].args, ("ch_xyz",)) + + async def test_join_raises_when_wrapper_returns_empty(self): + pool = _FakePool(fetchval_return=None) + backend = IsolatedPgmqBackend(pool=pool) + with self.assertRaises(RuntimeError): + await backend.join("ch_xyz") + + async def test_publish_serializes_payload_as_jsonb(self): + pool = _FakePool() + backend = IsolatedPgmqBackend(pool=pool) + + await backend.publish("ch_xyz", "q_self", {"k": "v", "n": 1}) + + self.assertEqual(len(pool.calls), 1) + call = pool.calls[0] + self.assertIn("public.bus_publish", call.sql) + # (channel, my_queue, json_text) + self.assertEqual(call.args[0], "ch_xyz") + self.assertEqual(call.args[1], "q_self") + self.assertEqual(json.loads(call.args[2]), {"k": "v", "n": 1}) + + async def test_read_normalizes_rows_and_decodes_string_jsonb(self): + # asyncpg can hand back jsonb as either a str (when no codec is + # registered) or a dict (with a codec). Both should normalize. + pool = _FakePool( + fetch_return=[ + {"msg_id": 1, "message": {"hello": "world"}}, + {"msg_id": 2, "message": json.dumps({"hello": "again"})}, + ] + ) + backend = IsolatedPgmqBackend(pool=pool) + + msgs = await backend.read( + "q_self", + channel="ch_xyz", + vt=30, + qty=10, + max_poll_seconds=5, + poll_interval_ms=100, + ) + + self.assertEqual(len(msgs), 2) + self.assertIsInstance(msgs[0], BackendMessage) + self.assertEqual(msgs[0].msg_id, 1) + self.assertEqual(msgs[0].message, {"hello": "world"}) + self.assertEqual(msgs[1].msg_id, 2) + self.assertEqual(msgs[1].message, {"hello": "again"}) + + self.assertEqual(len(pool.calls), 1) + call = pool.calls[0] + self.assertIn("public.bus_subscribe", call.sql) + # (my_queue, channel, vt, qty, max_poll_seconds) + self.assertEqual(call.args, ("q_self", "ch_xyz", 30, 10, 5)) + + async def test_archive_and_leave_call_the_right_wrappers(self): + pool = _FakePool(fetchval_return=True) + backend = IsolatedPgmqBackend(pool=pool) + + archived = await backend.archive("q_self", channel="ch_xyz", msg_id=42) + self.assertTrue(archived) + + await backend.leave("q_self", channel="ch_xyz") + + self.assertEqual(len(pool.calls), 2) + self.assertIn("public.bus_archive", pool.calls[0].sql) + self.assertEqual(pool.calls[0].args, ("q_self", "ch_xyz", 42)) + self.assertIn("public.bus_leave", pool.calls[1].sql) + self.assertEqual(pool.calls[1].args, ("q_self", "ch_xyz")) + + def test_satisfies_pgmq_backend_protocol(self): + # Runtime-checkable Protocol: confirm structural compliance. + backend = IsolatedPgmqBackend(pool=_FakePool()) + self.assertIsInstance(backend, PgmqBackend) + + +# --------------------------------------------------------------------------- +# PgmqBus(backend=...): orchestrator delegates to whatever backend it's given. +# --------------------------------------------------------------------------- + + +class _RecordingBackend: + """A backend that records every call PgmqBus makes against it.""" + + def __init__(self): + self.calls: list[tuple[str, tuple, dict]] = [] + self._inbound: asyncio.Queue[BackendMessage] = asyncio.Queue() + + async def join(self, channel): + self.calls.append(("join", (channel,), {})) + return "q_injected" + + async def publish(self, channel, my_queue, payload): + self.calls.append(("publish", (channel, my_queue, payload), {})) + + async def read(self, queue, *, channel, vt, qty, max_poll_seconds, poll_interval_ms): + self.calls.append( + ( + "read", + (queue,), + dict( + channel=channel, + vt=vt, + qty=qty, + max_poll_seconds=max_poll_seconds, + poll_interval_ms=poll_interval_ms, + ), + ) + ) + try: + msg = await asyncio.wait_for(self._inbound.get(), timeout=0.05) + return [msg] + except TimeoutError: + return [] + + async def archive(self, queue, *, channel, msg_id): + self.calls.append(("archive", (queue,), dict(channel=channel, msg_id=msg_id))) + return True + + async def leave(self, queue, *, channel): + self.calls.append(("leave", (queue,), dict(channel=channel))) + + +class TestPgmqBusBackendInjection(unittest.IsolatedAsyncioTestCase): + def test_rejects_construction_with_neither_pgmq_nor_backend(self): + with self.assertRaises(ValueError): + PgmqBus(channel="ch_xyz") + + def test_rejects_construction_with_both_pgmq_and_backend(self): + with self.assertRaises(ValueError): + PgmqBus(pgmq=object(), backend=_RecordingBackend(), channel="ch_xyz") + + async def test_lifecycle_and_publish_delegate_to_backend(self): + backend = _RecordingBackend() + bus = PgmqBus( + backend=backend, + channel="ch_xyz", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + + await bus.start() + try: + self.assertEqual(bus._queue_name, "q_injected") + await bus.send(BusDataMessage(source="task_a")) + # Give the (mostly idle) reader a chance to record a read call. + await asyncio.sleep(0.05) + finally: + await bus.stop() + + # join + publish (at least 1) + leave must all be observed; read may + # have fired multiple times, archive never (no inbound msg). + methods = [c[0] for c in backend.calls] + self.assertEqual(methods.count("join"), 1) + self.assertGreaterEqual(methods.count("publish"), 1) + self.assertEqual(methods.count("leave"), 1) + self.assertEqual(methods.count("archive"), 0) + + # The publish call passed (channel, my_queue, payload-dict). + publish_call = next(c for c in backend.calls if c[0] == "publish") + channel, my_queue, payload = publish_call[1] + self.assertEqual(channel, "ch_xyz") + self.assertEqual(my_queue, "q_injected") + self.assertIsInstance(payload, dict) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pgmq_bus.py b/tests/test_pgmq_bus.py new file mode 100644 index 000000000..b8c09d46f --- /dev/null +++ b/tests/test_pgmq_bus.py @@ -0,0 +1,611 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import itertools +import json +import unittest +from dataclasses import dataclass, field +from datetime import datetime + +from pipecat.bus import ( + BusAddTaskMessage, + BusDataMessage, + BusEndMessage, + BusFrameMessage, + BusJobRequestMessage, + BusSubscriber, +) +from pipecat.bus.network.pgmq import PgmqBus +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.frames.frames import TextFrame +from pipecat.pipeline.base_task import BaseTask +from pipecat.processors.frame_processor import FrameDirection +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + +_sub_counter = itertools.count() + + +def _make_sub(received: list) -> BusSubscriber: + """Create a BusSubscriber that appends messages to the given list.""" + sub_name = f"test_sub_{next(_sub_counter)}" + + class _Sub(BusSubscriber): + @property + def name(self) -> str: + return sub_name + + async def on_bus_message(self, message): + received.append(message) + + return _Sub() + + +@dataclass +class _FakeMessage: + msg_id: int + message: dict + read_ct: int = 0 + enqueued_at: datetime = field(default_factory=datetime.utcnow) + vt: datetime = field(default_factory=datetime.utcnow) + + +class FakePgmq: + """In-memory fake PGMQueue covering the surface used by ``PgmqBus``.""" + + def __init__(self): + self._queues: dict[str, list[_FakeMessage]] = {} + self._next_id = 1 + self.sent: list[tuple[str, dict]] = [] + + async def create_queue(self, queue: str, unlogged: bool = False) -> None: + self._queues.setdefault(queue, []) + + async def drop_queue(self, queue: str, partitioned: bool = False) -> bool: + existed = queue in self._queues + self._queues.pop(queue, None) + return existed + + async def list_queues(self) -> list[str]: + return list(self._queues.keys()) + + async def send(self, queue: str, message: dict, delay: int = 0, tz=None) -> int: + if queue not in self._queues: + raise RuntimeError(f"queue '{queue}' does not exist") + msg_id = self._next_id + self._next_id += 1 + self._queues[queue].append(_FakeMessage(msg_id=msg_id, message=message)) + self.sent.append((queue, message)) + return msg_id + + async def read_with_poll( + self, + queue: str, + vt: int | None = None, + qty: int = 1, + max_poll_seconds: int = 5, + poll_interval_ms: int = 100, + ) -> list[_FakeMessage]: + elapsed = 0.0 + interval = poll_interval_ms / 1000.0 + timeout = float(max_poll_seconds) + while elapsed < timeout: + queue_msgs = self._queues.get(queue) + if queue_msgs: + taken = queue_msgs[:qty] + # Mark as read by removing for now (simulates vt with a long + # enough vt that tests don't observe redelivery). + self._queues[queue] = queue_msgs[qty:] + return taken + await asyncio.sleep(interval) + elapsed += interval + return [] + + async def delete(self, queue: str, msg_id: int) -> bool: + # Messages in this fake are removed from the visible queue when read. + # delete() is essentially a no-op for tracking, but we record the call. + return True + + def inject(self, queue: str, message: dict) -> int: + """Push a message into a queue without going through send().""" + if queue not in self._queues: + self._queues[queue] = [] + msg_id = self._next_id + self._next_id += 1 + self._queues[queue].append(_FakeMessage(msg_id=msg_id, message=message)) + return msg_id + + +async def create_test_pgmq_bus( + channel: str = "test_bus", +) -> tuple[PgmqBus, FakePgmq, JSONMessageSerializer]: + """Create a PgmqBus with FakePgmq and TaskManager for testing.""" + pgmq = FakePgmq() + serializer = JSONMessageSerializer() + bus = PgmqBus( + pgmq=pgmq, + serializer=serializer, + channel=channel, + # Tight timing so tests run fast + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + return bus, pgmq, serializer + + +class TestPgmqBus(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.pgmq, self.serializer = await create_test_pgmq_bus() + await self.bus.start() + + async def asyncTearDown(self): + await self.bus.stop() + + async def test_send_publishes_to_pgmq(self): + """send() serializes and publishes to the PGMQ queue.""" + msg = BusDataMessage(source="task_a", target="task_b") + await self.bus.send(msg) + + self.assertEqual(len(self.pgmq.sent), 1) + queue, payload = self.pgmq.sent[0] + self.assertTrue(queue.startswith("test_bus_")) + self.assertIsInstance(payload, dict) + + restored = self.serializer.deserialize(json.dumps(payload).encode("utf-8")) + self.assertEqual(restored.source, "task_a") + self.assertEqual(restored.target, "task_b") + + async def test_local_message_stays_local(self): + """BusLocalMessage messages are delivered locally but not over PGMQ.""" + + received = [] + await self.bus.subscribe(_make_sub(received)) + + task = BaseTask("test") + msg = BusAddTaskMessage(source="parent", task=task) + await self.bus.send(msg) + + await asyncio.sleep(0.05) + + # Not published to PGMQ + self.assertEqual(len(self.pgmq.sent), 0) + # But delivered locally + self.assertEqual(len(received), 1) + self.assertIsInstance(received[0], BusAddTaskMessage) + self.assertIs(received[0].task, task) + + async def test_round_trip_via_subscriber(self): + """Messages published are received by subscribers.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + + msg = BusEndMessage(source="task_a", reason="done") + await self.bus.send(msg) + + await asyncio.sleep(0.2) + + self.assertEqual(len(received), 1) + self.assertIsInstance(received[0], BusEndMessage) + self.assertEqual(received[0].source, "task_a") + self.assertEqual(received[0].reason, "done") + + async def test_multiple_subscribers_receive(self): + """Multiple subscribers each receive every message.""" + received_a = [] + received_b = [] + await self.bus.subscribe(_make_sub(received_a)) + await self.bus.subscribe(_make_sub(received_b)) + + msg = BusDataMessage(source="x") + await self.bus.send(msg) + + await asyncio.sleep(0.2) + + self.assertEqual(len(received_a), 1) + self.assertEqual(len(received_b), 1) + + async def test_frame_message_round_trip(self): + """BusFrameMessage with a frame adapter round-trips through PGMQ.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + + msg = BusFrameMessage( + source="task_a", + frame=TextFrame(text="hello"), + direction=FrameDirection.DOWNSTREAM, + ) + await self.bus.send(msg) + + await asyncio.sleep(0.2) + + self.assertEqual(len(received), 1) + restored = received[0] + self.assertIsInstance(restored, BusFrameMessage) + self.assertIsInstance(restored.frame, TextFrame) + self.assertEqual(restored.frame.text, "hello") + self.assertEqual(restored.direction, FrameDirection.DOWNSTREAM) + + async def test_job_request_round_trip(self): + """BusJobRequestMessage round-trips through PGMQ.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + + msg = BusJobRequestMessage( + source="parent", + target="worker", + job_id="t-1", + payload={"key": "value"}, + ) + await self.bus.send(msg) + + await asyncio.sleep(0.2) + + self.assertEqual(len(received), 1) + restored = received[0] + self.assertIsInstance(restored, BusJobRequestMessage) + self.assertEqual(restored.job_id, "t-1") + self.assertEqual(restored.payload, {"key": "value"}) + + +class TestPgmqBusBroadcast(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.pgmq = FakePgmq() + self.serializer = JSONMessageSerializer() + self.tm = TaskManager() + self.tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + + self.bus_a = PgmqBus( + pgmq=self.pgmq, + serializer=self.serializer, + channel="test_bus", + poll_interval_ms=10, + max_poll_seconds=1, + ) + await self.bus_a.setup(self.tm) + + self.bus_b = PgmqBus( + pgmq=self.pgmq, + serializer=self.serializer, + channel="test_bus", + poll_interval_ms=10, + max_poll_seconds=1, + ) + await self.bus_b.setup(self.tm) + + async def asyncTearDown(self): + await self.bus_a.stop() + await self.bus_b.stop() + + async def test_broadcast_to_multiple_instances(self): + """Publishing on one instance reaches subscribers on a peer instance.""" + await self.bus_a.start() + await self.bus_b.start() + + received_a = [] + received_b = [] + await self.bus_a.subscribe(_make_sub(received_a)) + await self.bus_b.subscribe(_make_sub(received_b)) + + # bus_a publishes — both bus_a and bus_b should receive (round-trip). + msg = BusEndMessage(source="task_a", reason="hello peers") + await self.bus_a.send(msg) + + await asyncio.sleep(0.3) + + self.assertEqual(len(received_a), 1) + self.assertEqual(len(received_b), 1) + self.assertIsInstance(received_a[0], BusEndMessage) + self.assertIsInstance(received_b[0], BusEndMessage) + + async def test_publish_skips_dropped_queue(self): + """If a peer queue disappears mid-publish the broadcast keeps going.""" + await self.bus_a.start() + await self.bus_b.start() + + received_b = [] + await self.bus_b.subscribe(_make_sub(received_b)) + + # Drop bus_a's own queue out from under it (simulates a crashed peer). + await self.pgmq.drop_queue(self.bus_a._queue_name) + + msg = BusDataMessage(source="task_a") + # Should not raise even though one of the two peer queues is gone. + await self.bus_a.send(msg) + + await asyncio.sleep(0.3) + + # bus_b still received the message + self.assertEqual(len(received_b), 1) + + +class TestPgmqBusEdgeCases(unittest.IsolatedAsyncioTestCase): + async def test_custom_channel(self): + """Channel parameter is honored in the queue name.""" + pgmq = FakePgmq() + bus = PgmqBus( + pgmq=pgmq, + channel="custom_channel", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + self.assertTrue(bus._queue_name.startswith("custom_channel_")) + await bus.stop() + + async def test_channel_sanitization(self): + """Invalid channel characters are replaced with underscores.""" + pgmq = FakePgmq() + bus = PgmqBus( + pgmq=pgmq, + channel="pipecat:acme/v2", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + + # Queue name must match PGMQ's allowed identifier pattern. + self.assertRegex(bus._queue_name, r"^[A-Za-z_][A-Za-z0-9_]*$") + # And contain the sanitized prefix. + self.assertTrue(bus._queue_name.startswith("pipecat_acme_v2_")) + + await bus.stop() + + async def test_stop_cleans_up(self): + """stop() cancels the reader task and drops the queue.""" + pgmq = FakePgmq() + bus = PgmqBus(pgmq=pgmq, channel="cleanup_test", poll_interval_ms=10, max_poll_seconds=1) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + + await bus.start() + queue_name = bus._queue_name + self.assertIsNotNone(queue_name) + self.assertIsNotNone(bus._reader_task) + self.assertIn(queue_name, await pgmq.list_queues()) + + await bus.stop() + self.assertIsNone(bus._queue_name) + self.assertIsNone(bus._reader_task) + self.assertNotIn(queue_name, await pgmq.list_queues()) + + async def test_deserialize_failure_is_handled(self): + """A bad payload on the queue is logged and the reader keeps going.""" + pgmq = FakePgmq() + bus = PgmqBus( + pgmq=pgmq, + channel="poison_test", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + + received = [] + await bus.subscribe(_make_sub(received)) + await bus.start() + + # Inject a bogus payload directly into the bus's own queue. + pgmq.inject(bus._queue_name, {"not": "a real bus message"}) + + # Then send a real message through the normal path. + await bus.send(BusEndMessage(source="task_a", reason="recovered")) + + await asyncio.sleep(0.3) + + # The poison message should be silently dropped; the real one delivered. + real = [m for m in received if isinstance(m, BusEndMessage)] + self.assertEqual(len(real), 1) + self.assertEqual(real[0].reason, "recovered") + + await bus.stop() + + +class TestPgmqBusErrorPaths(unittest.IsolatedAsyncioTestCase): + """Cover error and edge branches in ``PgmqBus``.""" + + async def test_sanitize_channel_leading_digit(self): + """A channel starting with a digit gets a ``q_`` prefix.""" + pgmq = FakePgmq() + bus = PgmqBus(pgmq=pgmq, channel="123corp", poll_interval_ms=10, max_poll_seconds=1) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + try: + self.assertTrue(bus._queue_name.startswith("q_123corp_")) + finally: + await bus.stop() + + async def test_peer_cache_hit_on_rapid_publish(self): + """A second publish within the TTL reuses the cached peer list.""" + pgmq = FakePgmq() + bus = PgmqBus(pgmq=pgmq, channel="cache_hit", poll_interval_ms=10, max_poll_seconds=1) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + try: + list_calls = {"n": 0} + real_list_queues = pgmq.list_queues + + async def counting_list_queues(): + list_calls["n"] += 1 + return await real_list_queues() + + pgmq.list_queues = counting_list_queues + + await bus.send(BusDataMessage(source="a")) + first = list_calls["n"] + await bus.send(BusDataMessage(source="a")) + self.assertEqual(list_calls["n"], first, "second publish must hit cache") + finally: + await bus.stop() + + async def test_publish_skips_when_serializer_returns_invalid_json(self): + """If the serializer's bytes aren't decodable JSON, publish bails out.""" + + class BadSerializer(JSONMessageSerializer): + def serialize(self, message): + return b"not-valid-json" + + pgmq = FakePgmq() + bus = PgmqBus( + pgmq=pgmq, + serializer=BadSerializer(), + channel="bad_serializer", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + try: + await bus.send(BusDataMessage(source="a")) + # Nothing should land on the wire because json.loads failed. + self.assertEqual(len(pgmq.sent), 0) + finally: + await bus.stop() + + async def test_publish_handles_send_failure_to_cached_peer(self): + """A queue dropped after the cache warmed lets publish keep going.""" + pgmq = FakePgmq() + bus = PgmqBus( + pgmq=pgmq, + channel="cached_drop", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + try: + # Warm the peer cache with the bus's own queue. + await bus.send(BusDataMessage(source="warmup")) + self.assertIn("cached_drop", bus._backend._peer_cache) + # Drop our queue so the cached entry now points at nothing. + await pgmq.drop_queue(bus._queue_name) + # Force-keep the cache fresh so we hit the cached path on the next publish. + cached_at, peers = bus._backend._peer_cache["cached_drop"] + bus._backend._peer_cache["cached_drop"] = (float("inf"), peers) + # Should swallow the exception from the dropped peer and reset the cache. + await bus.send(BusDataMessage(source="dead-letter")) + self.assertNotIn("cached_drop", bus._backend._peer_cache) + finally: + # Re-create the queue so stop() can drop it cleanly without surfacing. + await pgmq.create_queue(bus._queue_name) + await bus.stop() + + async def test_stop_swallows_drop_queue_failure(self): + """``stop()`` logs but doesn't raise when ``drop_queue`` fails.""" + + class BadDropPgmq(FakePgmq): + async def drop_queue(self, queue, partitioned=False): + raise RuntimeError("simulated drop failure") + + pgmq = BadDropPgmq() + bus = PgmqBus(pgmq=pgmq, channel="bad_drop", poll_interval_ms=10, max_poll_seconds=1) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + + # Should not raise even though drop_queue throws. + await bus.stop() + self.assertIsNone(bus._queue_name) + + async def test_reader_loop_recovers_from_read_failure(self): + """A failure inside ``read_with_poll`` triggers backoff, not crash.""" + + class FlakyReadPgmq(FakePgmq): + def __init__(self): + super().__init__() + self.read_calls = 0 + + async def read_with_poll(self, *args, **kwargs): + self.read_calls += 1 + if self.read_calls == 1: + raise RuntimeError("simulated read failure") + return await super().read_with_poll(*args, **kwargs) + + pgmq = FlakyReadPgmq() + bus = PgmqBus(pgmq=pgmq, channel="flaky_read", poll_interval_ms=10, max_poll_seconds=1) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + try: + # Give the reader long enough to fail once and back off. + await asyncio.sleep(1.2) + self.assertGreaterEqual(pgmq.read_calls, 2) + finally: + await bus.stop() + + async def test_reader_loop_handles_deserializer_exception(self): + """If the serializer raises while decoding, the message is still deleted.""" + + class RaisingSerializer(JSONMessageSerializer): + def deserialize(self, data): + raise RuntimeError("simulated deserialize crash") + + pgmq = FakePgmq() + bus = PgmqBus( + pgmq=pgmq, + serializer=RaisingSerializer(), + channel="bad_deserialize", + poll_interval_ms=10, + max_poll_seconds=1, + ) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + await bus.start() + try: + pgmq.inject(bus._queue_name, {"junk": True}) + # Reader should consume + log + continue without crashing. + await asyncio.sleep(0.3) + finally: + await bus.stop() + + async def test_reader_loop_logs_when_delete_fails(self): + """A failing ``delete`` after dispatch is logged and not propagated.""" + + class BadDeletePgmq(FakePgmq): + async def delete(self, queue, msg_id): + raise RuntimeError("simulated delete failure") + + pgmq = BadDeletePgmq() + bus = PgmqBus(pgmq=pgmq, channel="bad_delete", poll_interval_ms=10, max_poll_seconds=1) + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + + received = [] + await bus.subscribe(_make_sub(received)) + await bus.start() + try: + await bus.send(BusEndMessage(source="a", reason="ok")) + await asyncio.sleep(0.3) + # Real message still delivered locally despite delete failing. + self.assertEqual(len(received), 1) + finally: + await bus.stop() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_redis_bus.py b/tests/test_redis_bus.py new file mode 100644 index 000000000..0f1cc9511 --- /dev/null +++ b/tests/test_redis_bus.py @@ -0,0 +1,249 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import itertools +import unittest + +from pipecat.bus import ( + BusAddTaskMessage, + BusDataMessage, + BusEndMessage, + BusFrameMessage, + BusJobRequestMessage, + BusSubscriber, +) +from pipecat.bus.network.redis import RedisBus +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.frames.frames import TextFrame +from pipecat.pipeline.base_task import BaseTask +from pipecat.processors.frame_processor import FrameDirection +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + +_sub_counter = itertools.count() + + +def _make_sub(received: list) -> BusSubscriber: + """Create a BusSubscriber that appends messages to the given list.""" + sub_name = f"test_sub_{next(_sub_counter)}" + + class _Sub(BusSubscriber): + @property + def name(self) -> str: + return sub_name + + async def on_bus_message(self, message): + received.append(message) + + return _Sub() + + +class FakePubSub: + """In-memory fake Redis pub/sub for testing.""" + + def __init__(self): + self._subscriptions: dict[str, asyncio.Queue] = {} + self._closed = False + + async def subscribe(self, channel: str): + self._subscriptions[channel] = asyncio.Queue() + + async def unsubscribe(self, channel: str): + self._subscriptions.pop(channel, None) + + async def close(self): + self._closed = True + + async def listen(self): + """Yield messages from the subscription queue.""" + # We only support one channel in tests + channel = next(iter(self._subscriptions)) + queue = self._subscriptions[channel] + while True: + msg = await queue.get() + yield msg + + def inject(self, channel: str, data: bytes): + """Inject a raw message into the fake pub/sub.""" + if channel in self._subscriptions: + self._subscriptions[channel].put_nowait( + {"type": "message", "data": data, "channel": channel} + ) + + +class FakeRedis: + """In-memory fake Redis client for testing.""" + + def __init__(self): + self._pubsubs: list[FakePubSub] = [] + self._published: list[tuple[str, bytes]] = [] + + def pubsub(self): + ps = FakePubSub() + self._pubsubs.append(ps) + return ps + + async def publish(self, channel: str, data: bytes): + self._published.append((channel, data)) + # Fan out to all pubsub instances + for ps in self._pubsubs: + ps.inject(channel, data) + + +async def create_test_redis_bus(): + """Create a RedisBus with fake Redis and TaskManager for testing.""" + redis = FakeRedis() + serializer = JSONMessageSerializer() + bus = RedisBus(redis=redis, serializer=serializer, channel="test:bus") + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + return bus, redis, serializer + + +class TestRedisBus(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.redis, self.serializer = await create_test_redis_bus() + + async def test_send_publishes_to_redis(self): + """send() serializes and publishes to the Redis channel.""" + msg = BusDataMessage(source="task_a", target="task_b") + await self.bus.send(msg) + + self.assertEqual(len(self.redis._published), 1) + channel, data = self.redis._published[0] + self.assertEqual(channel, "test:bus") + self.assertIsInstance(data, bytes) + + # Verify it deserializes back + restored = self.serializer.deserialize(data) + self.assertEqual(restored.source, "task_a") + self.assertEqual(restored.target, "task_b") + + async def test_local_mixin_delivered_locally_not_to_redis(self): + """BusLocalMessage messages are delivered to local subscribers but not published to Redis.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + await self.bus.start() + + task = BaseTask("test") + msg = BusAddTaskMessage(source="parent", task=task) + await self.bus.send(msg) + + await asyncio.sleep(0.05) + await self.bus.stop() + + # Not published to Redis + self.assertEqual(len(self.redis._published), 0) + # But delivered locally + self.assertEqual(len(received), 1) + self.assertIsInstance(received[0], BusAddTaskMessage) + self.assertIs(received[0].task, task) + + async def test_round_trip_via_subscriber(self): + """Messages published are received by subscribers.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + await self.bus.start() + + msg = BusEndMessage(source="task_a", reason="done") + await self.bus.send(msg) + + # Give the reader task time to process + await asyncio.sleep(0.1) + await self.bus.stop() + + self.assertEqual(len(received), 1) + self.assertIsInstance(received[0], BusEndMessage) + self.assertEqual(received[0].source, "task_a") + self.assertEqual(received[0].reason, "done") + + async def test_multiple_subscribers_receive(self): + """Multiple subscribers each receive every message.""" + received_a = [] + received_b = [] + await self.bus.subscribe(_make_sub(received_a)) + await self.bus.subscribe(_make_sub(received_b)) + await self.bus.start() + + msg = BusDataMessage(source="x") + await self.bus.send(msg) + + await asyncio.sleep(0.1) + await self.bus.stop() + + self.assertEqual(len(received_a), 1) + self.assertEqual(len(received_b), 1) + + async def test_frame_message_round_trip(self): + """BusFrameMessage with a frame adapter round-trips through Redis.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + await self.bus.start() + + msg = BusFrameMessage( + source="task_a", + frame=TextFrame(text="hello"), + direction=FrameDirection.DOWNSTREAM, + ) + await self.bus.send(msg) + + await asyncio.sleep(0.1) + await self.bus.stop() + + self.assertEqual(len(received), 1) + restored = received[0] + self.assertIsInstance(restored, BusFrameMessage) + self.assertIsInstance(restored.frame, TextFrame) + self.assertEqual(restored.frame.text, "hello") + self.assertEqual(restored.direction, FrameDirection.DOWNSTREAM) + + async def test_job_request_round_trip(self): + """BusJobRequestMessage round-trips through Redis.""" + received = [] + await self.bus.subscribe(_make_sub(received)) + await self.bus.start() + + msg = BusJobRequestMessage( + source="parent", + target="worker", + job_id="t-1", + payload={"key": "value"}, + ) + await self.bus.send(msg) + + await asyncio.sleep(0.1) + await self.bus.stop() + + self.assertEqual(len(received), 1) + restored = received[0] + self.assertIsInstance(restored, BusJobRequestMessage) + self.assertEqual(restored.job_id, "t-1") + self.assertEqual(restored.payload, {"key": "value"}) + + async def test_custom_channel(self): + """Messages are published to the configured channel.""" + bus = RedisBus( + redis=self.redis, + serializer=self.serializer, + channel="custom:channel", + ) + await bus.send(BusDataMessage(source="a")) + self.assertEqual(self.redis._published[0][0], "custom:channel") + + async def test_stop_cleans_up(self): + """stop() cancels the reader task and unsubscribes from Redis.""" + await self.bus.start() + self.assertIsNotNone(self.bus._pubsub) + self.assertIsNotNone(self.bus._reader_task) + + await self.bus.stop() + self.assertIsNone(self.bus._pubsub) + self.assertIsNone(self.bus._reader_task) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_registry.py b/tests/test_registry.py new file mode 100644 index 000000000..1982ab8e5 --- /dev/null +++ b/tests/test_registry.py @@ -0,0 +1,170 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import unittest + +from pipecat.registry import TaskRegistry +from pipecat.registry.types import TaskReadyData + + +class TestTaskRegistry(unittest.IsolatedAsyncioTestCase): + def setUp(self): + self.registry = TaskRegistry(runner_name="runner_a") + + async def test_register_local_task(self): + """Local task is registered and appears in local_tasks.""" + data = TaskReadyData(task_name="greeter", runner="runner_a") + result = await self.registry.register(data) + + self.assertTrue(result) + self.assertIn("greeter", self.registry.local_tasks) + self.assertNotIn("greeter", self.registry.remote_tasks) + + async def test_register_remote_task(self): + """Remote task is registered and appears in remote_tasks.""" + data = TaskReadyData(task_name="support", runner="runner_b") + result = await self.registry.register(data) + + self.assertTrue(result) + self.assertIn("support", self.registry.remote_tasks) + self.assertNotIn("support", self.registry.local_tasks) + + async def test_duplicate_registration_returns_false(self): + """Registering the same task twice returns False.""" + data = TaskReadyData(task_name="greeter", runner="runner_a") + first = await self.registry.register(data) + second = await self.registry.register(data) + + self.assertTrue(first) + self.assertFalse(second) + self.assertEqual(self.registry.local_tasks.count("greeter"), 1) + + async def test_get_local_task(self): + """get() returns data for a local task.""" + data = TaskReadyData(task_name="greeter", runner="runner_a") + await self.registry.register(data) + + result = self.registry.get("greeter") + self.assertIs(result, data) + + async def test_get_remote_task(self): + """get() returns data for a remote task.""" + data = TaskReadyData(task_name="support", runner="runner_b") + await self.registry.register(data) + + result = self.registry.get("support") + self.assertIs(result, data) + + async def test_get_unknown_task_returns_none(self): + """get() returns None for an unknown task.""" + self.assertIsNone(self.registry.get("nonexistent")) + + async def test_contains(self): + """__contains__ works for registered and unregistered tasks.""" + data = TaskReadyData(task_name="greeter", runner="runner_a") + await self.registry.register(data) + + self.assertIn("greeter", self.registry) + self.assertNotIn("unknown", self.registry) + + async def test_watch_fires_on_registration(self): + """Watch handler fires when the watched task registers.""" + received = [] + + async def handler(task_data): + received.append(task_data) + + await self.registry.watch("greeter", handler) + + data = TaskReadyData(task_name="greeter", runner="runner_a") + await self.registry.register(data) + + self.assertEqual(len(received), 1) + self.assertIs(received[0], data) + + async def test_watch_does_not_fire_for_other_tasks(self): + """Watch handler does not fire for a different task.""" + received = [] + + async def handler(task_data): + received.append(task_data) + + await self.registry.watch("greeter", handler) + + data = TaskReadyData(task_name="support", runner="runner_a") + await self.registry.register(data) + + self.assertEqual(len(received), 0) + + async def test_watch_does_not_fire_on_duplicate(self): + """Watch handler does not fire on duplicate registration.""" + received = [] + + async def handler(task_data): + received.append(task_data) + + await self.registry.watch("greeter", handler) + + data = TaskReadyData(task_name="greeter", runner="runner_a") + await self.registry.register(data) + await self.registry.register(data) + + self.assertEqual(len(received), 1) + + async def test_multiple_watchers(self): + """Multiple watch handlers fire for the same task.""" + received_a = [] + received_b = [] + + async def handler_a(task_data): + received_a.append(task_data) + + async def handler_b(task_data): + received_b.append(task_data) + + await self.registry.watch("greeter", handler_a) + await self.registry.watch("greeter", handler_b) + + data = TaskReadyData(task_name="greeter", runner="runner_a") + await self.registry.register(data) + + self.assertEqual(len(received_a), 1) + self.assertEqual(len(received_b), 1) + + async def test_watch_fires_immediately_if_already_registered(self): + """Watch handler fires immediately when the task is already registered.""" + data = TaskReadyData(task_name="greeter", runner="runner_a") + await self.registry.register(data) + + received = [] + + async def handler(task_data): + received.append(task_data) + + await self.registry.watch("greeter", handler) + + self.assertEqual(len(received), 1) + self.assertIs(received[0], data) + + async def test_runner_name_property(self): + """runner_name returns the name passed at construction.""" + self.assertEqual(self.registry.runner_name, "runner_a") + + async def test_multiple_remote_runners(self): + """Tasks from multiple remote runners are tracked separately.""" + data_b = TaskReadyData(task_name="task_b", runner="runner_b") + data_c = TaskReadyData(task_name="task_c", runner="runner_c") + await self.registry.register(data_b) + await self.registry.register(data_c) + + remote = self.registry.remote_tasks + self.assertIn("task_b", remote) + self.assertIn("task_c", remote) + self.assertEqual(len(remote), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_runner.py b/tests/test_runner.py new file mode 100644 index 000000000..67a6560c4 --- /dev/null +++ b/tests/test_runner.py @@ -0,0 +1,202 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest + +from pipecat.bus import ( + BusAddTaskMessage, + BusCancelMessage, + BusCancelTaskMessage, + BusEndMessage, + BusEndTaskMessage, +) +from pipecat.pipeline.base_task import BaseTask +from pipecat.pipeline.runner import PipelineRunner + + +class StubTask(BaseTask): + """BaseTask subclass that stops on end/cancel so the runner can exit.""" + + async def _handle_task_end(self, message): + await super()._handle_task_end(message) + self._finished_event.set() + + async def _handle_task_cancel(self, message): + await super()._handle_task_cancel(message) + self._finished_event.set() + + +class TestPipelineRunner(unittest.IsolatedAsyncioTestCase): + async def test_spawn_registers_task(self): + """spawn() registers the task by name (duplicate is silently skipped).""" + runner = PipelineRunner(handle_sigint=False) + task = StubTask("task_a") + + await runner.spawn(task) + + # Duplicate is silently skipped (logs error) + await runner.spawn(StubTask("task_a")) + + async def test_run_starts_bus_and_tasks(self): + """run() starts bus, starts all tasks, fires on_ready.""" + runner = PipelineRunner(handle_sigint=False) + task = StubTask("task_a") + await runner.spawn(task) + + runner_started = asyncio.Event() + + @runner.event_handler("on_ready") + async def on_ready(runner): + runner_started.set() + # Immediately end to unblock run() + await runner.end() + + await asyncio.wait_for(runner.run(), timeout=5.0) + + self.assertTrue(runner_started.is_set()) + + async def test_end_is_idempotent(self): + """end() is idempotent — subsequent calls are no-ops.""" + runner = PipelineRunner(handle_sigint=False) + task = StubTask("task_a") + await runner.spawn(task) + + @runner.event_handler("on_ready") + async def on_ready(runner): + await runner.end(reason="first") + await runner.end(reason="second") # should be no-op + + await asyncio.wait_for(runner.run(), timeout=5.0) + # If we got here without hanging, idempotency works + + async def test_cancel_is_idempotent(self): + """cancel() is idempotent — subsequent calls are no-ops.""" + runner = PipelineRunner(handle_sigint=False) + task = StubTask("task_a") + await runner.spawn(task) + + @runner.event_handler("on_ready") + async def on_ready(runner): + await runner.cancel(reason="first") + await runner.cancel(reason="second") # should be no-op + + try: + await asyncio.wait_for(runner.run(), timeout=5.0) + except asyncio.CancelledError: + pass + + async def test_end_sends_end_task_message_to_root_tasks_only(self): + """end() sends BusEndTaskMessage only to root tasks (no parent).""" + runner = PipelineRunner(handle_sigint=False) + root = StubTask("root") + child = StubTask("child") + # Manually mark child as having root as parent + child._parent = root.name + await runner.spawn(root) + await runner.spawn(child) + + sent = [] + bus = runner.bus + original_send = bus.send + + async def capture_send(message): + sent.append(message) + await original_send(message) + + bus.send = capture_send + + # Call end() directly — no need to run the full pipeline lifecycle + await runner.end() + + end_msgs = [m for m in sent if isinstance(m, BusEndTaskMessage)] + targets = {m.target for m in end_msgs} + self.assertIn("root", targets) + self.assertNotIn("child", targets) + + async def test_cancel_sends_cancel_task_message_to_root_tasks_only(self): + """cancel() sends BusCancelTaskMessage only to root tasks (no parent).""" + runner = PipelineRunner(handle_sigint=False) + root = StubTask("root") + child = StubTask("child") + child._parent = root.name + await runner.spawn(root) + await runner.spawn(child) + + sent = [] + bus = runner.bus + original_send = bus.send + + async def capture_send(message): + sent.append(message) + await original_send(message) + + bus.send = capture_send + + # Call cancel() directly — no need to run the full pipeline lifecycle + await runner.cancel() + + cancel_msgs = [m for m in sent if isinstance(m, BusCancelTaskMessage)] + targets = {m.target for m in cancel_msgs} + self.assertIn("root", targets) + self.assertNotIn("child", targets) + + async def test_bus_end_message_triggers_end(self): + """BusEndMessage on bus triggers runner.end().""" + runner = PipelineRunner(handle_sigint=False) + task = StubTask("task_a") + await runner.spawn(task) + + bus = runner.bus + + @runner.event_handler("on_ready") + async def on_ready(runner): + # Simulate a task sending BusEndMessage + await bus.send(BusEndMessage(source="task_a")) + + await asyncio.wait_for(runner.run(), timeout=5.0) + # If we got here, end was triggered by the bus message + + async def test_bus_cancel_message_triggers_cancel(self): + """BusCancelMessage on bus triggers runner.cancel().""" + runner = PipelineRunner(handle_sigint=False) + task = StubTask("task_a") + await runner.spawn(task) + + bus = runner.bus + + @runner.event_handler("on_ready") + async def on_ready(runner): + await bus.send(BusCancelMessage(source="task_a")) + + try: + await asyncio.wait_for(runner.run(), timeout=5.0) + except asyncio.CancelledError: + pass + + async def test_bus_add_task_message_triggers_spawn(self): + """BusAddTaskMessage on bus triggers spawn().""" + runner = PipelineRunner(handle_sigint=False) + task_a = StubTask("task_a") + await runner.spawn(task_a) + + task_b = StubTask("task_b") + bus = runner.bus + + @runner.event_handler("on_ready") + async def on_ready(runner): + await bus.send(BusAddTaskMessage(source="task_a", task=task_b)) + await asyncio.sleep(0.1) + await runner.end() + + await asyncio.wait_for(runner.run(), timeout=5.0) + + # Verify task_b was added (duplicate is silently skipped) + await runner.spawn(StubTask("task_b")) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_serializers.py b/tests/test_serializers.py new file mode 100644 index 000000000..d3ef97ad0 --- /dev/null +++ b/tests/test_serializers.py @@ -0,0 +1,246 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import unittest + +from pydantic import BaseModel + +from pipecat.bus.messages import ( + BusActivateTaskMessage, + BusCancelMessage, + BusDataMessage, + BusEndMessage, + BusFrameMessage, + BusJobRequestMessage, + BusJobResponseMessage, + BusMessage, +) +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.frames.frames import TextFrame +from pipecat.pipeline.job_context import JobStatus +from pipecat.processors.frame_processor import FrameDirection + + +class _Address(BaseModel): + city: str + zip_code: str + + +class _UserInfo(BaseModel): + name: str + age: int + address: _Address | None = None + + +class TestJSONMessageSerializer(unittest.TestCase): + def setUp(self): + self.serializer = JSONMessageSerializer() + + def test_round_trip_simple_message(self): + """BusMessage serializes and deserializes correctly.""" + msg = BusDataMessage(source="task_a", target="task_b") + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusMessage) + self.assertEqual(restored.source, "task_a") + self.assertEqual(restored.target, "task_b") + + def test_round_trip_broadcast_message(self): + """Broadcast message (no target) round-trips.""" + msg = BusDataMessage(source="task_a") + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusMessage) + self.assertEqual(restored.source, "task_a") + self.assertIsNone(restored.target) + + def test_round_trip_activate_message(self): + """BusActivateTaskMessage with args round-trips.""" + msg = BusActivateTaskMessage( + source="parent", + target="child", + args={"messages": [{"role": "user", "content": "hello"}]}, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusActivateTaskMessage) + self.assertEqual(restored.source, "parent") + self.assertEqual(restored.target, "child") + self.assertEqual(restored.args["messages"][0]["content"], "hello") + + def test_round_trip_end_message(self): + """BusEndMessage round-trips.""" + msg = BusEndMessage(source="task_a", reason="done") + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusEndMessage) + self.assertEqual(restored.reason, "done") + + def test_round_trip_cancel_message(self): + """BusCancelMessage round-trips.""" + msg = BusCancelMessage(source="task_a", reason="abort") + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusCancelMessage) + self.assertEqual(restored.reason, "abort") + + def test_round_trip_job_request(self): + """BusJobRequestMessage with payload round-trips.""" + msg = BusJobRequestMessage( + source="parent", + target="worker", + job_id="t-123", + payload={"key": "value"}, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusJobRequestMessage) + self.assertEqual(restored.job_id, "t-123") + self.assertEqual(restored.payload, {"key": "value"}) + + def test_round_trip_job_response(self): + """BusJobResponseMessage round-trips.""" + msg = BusJobResponseMessage( + source="worker", + target="parent", + job_id="t-123", + status=JobStatus.COMPLETED, + response={"result": 42}, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusJobResponseMessage) + self.assertEqual(restored.job_id, "t-123") + self.assertEqual(restored.response, {"result": 42}) + self.assertEqual(restored.status, JobStatus.COMPLETED) + + def test_round_trip_frame_message(self): + """BusFrameMessage with TextFrame round-trips via adapter.""" + msg = BusFrameMessage( + source="task_a", + frame=TextFrame(text="hello world"), + direction=FrameDirection.DOWNSTREAM, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusFrameMessage) + self.assertIsInstance(restored.frame, TextFrame) + self.assertEqual(restored.frame.text, "hello world") + self.assertEqual(restored.direction, FrameDirection.DOWNSTREAM) + self.assertEqual(restored.source, "task_a") + + def test_frame_message_upstream_direction(self): + """UPSTREAM direction preserved in round-trip.""" + msg = BusFrameMessage( + source="task_a", + frame=TextFrame(text="up"), + direction=FrameDirection.UPSTREAM, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertEqual(restored.direction, FrameDirection.UPSTREAM) + + def test_unregistered_frame_warns_and_skips(self): + """Serializing a frame with no adapter warns and skips the field.""" + serializer = JSONMessageSerializer() # no adapters registered + + msg = BusFrameMessage( + source="task_a", + frame=TextFrame(text="hello"), + direction=FrameDirection.DOWNSTREAM, + ) + # Should not raise — unserializable field is skipped with a warning + data = serializer.serialize(msg) + self.assertIsInstance(data, bytes) + + def test_unknown_message_type_returns_none(self): + """Deserializing an unknown message type returns None.""" + bad_data = b'{"__type__":"bogus.BogusMessage","__data__":{"source":"a"}}' + result = self.serializer.deserialize(bad_data) + self.assertIsNone(result) + + def test_serialized_is_bytes(self): + """serialize() returns bytes.""" + msg = BusDataMessage(source="a") + data = self.serializer.serialize(msg) + self.assertIsInstance(data, bytes) + + def test_adapter_mro_lookup(self): + """Adapter registered for a parent class handles subclasses.""" + + class CustomTextFrame(TextFrame): + pass + + msg = BusFrameMessage( + source="a", + frame=CustomTextFrame(text="sub"), + direction=FrameDirection.DOWNSTREAM, + ) + # TextTypeAdapter is registered for TextFrame, should handle subclass + data = self.serializer.serialize(msg) + self.assertIsInstance(data, bytes) + + def test_round_trip_pydantic_base_model(self): + """Pydantic BaseModel round-trips preserving the type.""" + msg = BusJobResponseMessage( + source="worker", + target="parent", + job_id="t-456", + status=JobStatus.COMPLETED, + response={"user": _UserInfo(name="Alice", age=30)}, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusJobResponseMessage) + user = restored.response["user"] + self.assertIsInstance(user, _UserInfo) + self.assertEqual(user.name, "Alice") + self.assertEqual(user.age, 30) + + def test_round_trip_nested_pydantic_base_model(self): + """Nested Pydantic BaseModels round-trip preserving types.""" + user = _UserInfo(name="Alice", age=30, address=_Address(city="NYC", zip_code="10001")) + msg = BusJobResponseMessage( + source="worker", + target="parent", + job_id="t-789", + status=JobStatus.COMPLETED, + response={"user": user}, + ) + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertIsInstance(restored, BusJobResponseMessage) + restored_user = restored.response["user"] + self.assertIsInstance(restored_user, _UserInfo) + self.assertIsInstance(restored_user.address, _Address) + self.assertEqual(restored_user.address.city, "NYC") + self.assertEqual(restored_user.address.zip_code, "10001") + + def test_non_init_fields_preserved(self): + """Non-init dataclass fields survive round-trip via setattr.""" + msg = BusDataMessage(source="task_a", target="task_b") + # name is a non-init field on DataFrame + msg.name = "custom_name" + + data = self.serializer.serialize(msg) + restored = self.serializer.deserialize(data) + + self.assertEqual(restored.name, "custom_name") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_websocket_proxy.py b/tests/test_websocket_proxy.py new file mode 100644 index 000000000..4bbd82f53 --- /dev/null +++ b/tests/test_websocket_proxy.py @@ -0,0 +1,318 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +import unittest +from unittest.mock import MagicMock + +from pipecat.bus import ( + AsyncQueueBus, + BusAddTaskMessage, + BusDataMessage, +) +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.pipeline.base_task import BaseTask +from pipecat.registry import TaskRegistry +from pipecat.utils.asyncio.task_manager import TaskManager, TaskManagerParams + + +async def create_test_bus(): + """Create an AsyncQueueBus with a TaskManager for testing.""" + bus = AsyncQueueBus() + tm = TaskManager() + tm.setup(TaskManagerParams(loop=asyncio.get_running_loop())) + await bus.setup(tm) + return bus, tm + + +class FakeWebSocket: + """Fake WebSocket for testing the client proxy.""" + + def __init__(self): + self._sent: list[bytes] = [] + self._receive_queue: asyncio.Queue[bytes] = asyncio.Queue() + self.closed = False + + async def send(self, data: bytes): + self._sent.append(data) + + async def close(self): + self.closed = True + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return await asyncio.wait_for(self._receive_queue.get(), timeout=0.5) + except TimeoutError: + raise StopAsyncIteration + + def inject(self, data: bytes): + """Inject data as if received from the remote side.""" + self._receive_queue.put_nowait(data) + + +class FakeStarletteWebSocket: + """Fake Starlette WebSocket for testing the server proxy.""" + + def __init__(self): + self._sent: list[bytes] = [] + self._receive_queue: asyncio.Queue[bytes] = asyncio.Queue() + self.closed = False + + async def send_bytes(self, data: bytes): + self._sent.append(data) + + async def receive_bytes(self) -> bytes: + return await self._receive_queue.get() + + async def close(self): + self.closed = True + + @property + def client_state(self): + return MagicMock(value=1) if not self.closed else MagicMock(value=3) + + def inject(self, data: bytes): + """Inject data as if received from the remote client.""" + self._receive_queue.put_nowait(data) + + +class TestWebSocketProxyClientTask(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm = await create_test_bus() + self.registry = TaskRegistry(runner_name="test-runner") + self.serializer = JSONMessageSerializer() + + async def _create_client(self, fake_ws): + from pipecat.tasks.proxy.websocket.client import WebSocketProxyClientTask + + task = WebSocketProxyClientTask( + "proxy", + url="ws://fake", + remote_task_name="worker", + local_task_name="voice", + serializer=self.serializer, + ) + task.attach(registry=self.registry, bus=self.bus) + await task.setup(self.tm) + task._ws = fake_ws + return task + + async def test_forwards_targeted_messages(self): + """Messages targeted at the remote task are forwarded.""" + fake_ws = FakeWebSocket() + task = await self._create_client(fake_ws) + + msg = BusDataMessage(source="voice", target="worker") + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 1) + restored = self.serializer.deserialize(fake_ws._sent[0]) + self.assertEqual(restored.source, "voice") + self.assertEqual(restored.target, "worker") + + async def test_skips_messages_for_other_tasks(self): + """Messages targeted at other tasks are not forwarded.""" + fake_ws = FakeWebSocket() + task = await self._create_client(fake_ws) + + msg = BusDataMessage(source="voice", target="other_task") + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 0) + + async def test_skips_broadcast_messages(self): + """Broadcast messages (target=None) are not forwarded.""" + fake_ws = FakeWebSocket() + task = await self._create_client(fake_ws) + + msg = BusDataMessage(source="voice") + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 0) + + async def test_skips_local_messages(self): + """BusLocalMessage messages are not forwarded.""" + fake_ws = FakeWebSocket() + task = await self._create_client(fake_ws) + + stub = MagicMock(spec=BaseTask) + stub.name = "child" + msg = BusAddTaskMessage(source="parent", target="worker", task=stub) + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 0) + + async def test_accepts_inbound_for_local_task(self): + """Inbound messages targeted at the local task are accepted.""" + fake_ws = FakeWebSocket() + task = await self._create_client(fake_ws) + + sent_to_bus = [] + original_send = task.send_message + + async def capture_send(message): + sent_to_bus.append(message) + await original_send(message) + + task.send_message = capture_send + + msg = BusDataMessage(source="worker", target="voice") + fake_ws.inject(self.serializer.serialize(msg)) + + receive_task = asyncio.create_task(task._receive_loop()) + await asyncio.sleep(0.1) + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + self.assertEqual(len(sent_to_bus), 1) + self.assertEqual(sent_to_bus[0].source, "worker") + self.assertEqual(sent_to_bus[0].target, "voice") + + async def test_drops_inbound_for_other_tasks(self): + """Inbound messages targeted at other tasks are dropped.""" + fake_ws = FakeWebSocket() + task = await self._create_client(fake_ws) + + sent_to_bus = [] + original_send = task.send_message + + async def capture_send(message): + sent_to_bus.append(message) + await original_send(message) + + task.send_message = capture_send + + msg = BusDataMessage(source="worker", target="other_task") + fake_ws.inject(self.serializer.serialize(msg)) + + receive_task = asyncio.create_task(task._receive_loop()) + await asyncio.sleep(0.1) + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + self.assertEqual(len(sent_to_bus), 0) + + +class TestWebSocketProxyServerTask(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.bus, self.tm = await create_test_bus() + self.registry = TaskRegistry(runner_name="test-runner") + self.serializer = JSONMessageSerializer() + + async def _create_server(self, fake_ws): + from pipecat.tasks.proxy.websocket.server import WebSocketProxyServerTask + + task = WebSocketProxyServerTask( + "gateway", + websocket=fake_ws, + task_name="worker", + remote_task_name="voice", + serializer=self.serializer, + ) + task.attach(registry=self.registry, bus=self.bus) + await task.setup(self.tm) + return task + + async def test_forwards_messages_from_local_task(self): + """Messages from the local task targeted at the remote task are forwarded.""" + fake_ws = FakeStarletteWebSocket() + task = await self._create_server(fake_ws) + + msg = BusDataMessage(source="worker", target="voice") + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 1) + restored = self.serializer.deserialize(fake_ws._sent[0]) + self.assertEqual(restored.source, "worker") + self.assertEqual(restored.target, "voice") + + async def test_skips_messages_from_other_tasks(self): + """Messages from other tasks are not forwarded.""" + fake_ws = FakeStarletteWebSocket() + task = await self._create_server(fake_ws) + + msg = BusDataMessage(source="other_task", target="voice") + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 0) + + async def test_skips_messages_to_other_targets(self): + """Messages from the local task to other targets are not forwarded.""" + fake_ws = FakeStarletteWebSocket() + task = await self._create_server(fake_ws) + + msg = BusDataMessage(source="worker", target="other_task") + await task.on_bus_message(msg) + + self.assertEqual(len(fake_ws._sent), 0) + + async def test_accepts_inbound_for_local_task(self): + """Inbound messages targeted at the local task are accepted.""" + fake_ws = FakeStarletteWebSocket() + task = await self._create_server(fake_ws) + + sent_to_bus = [] + original_send = task.send_message + + async def capture_send(message): + sent_to_bus.append(message) + await original_send(message) + + task.send_message = capture_send + + msg = BusDataMessage(source="voice", target="worker") + fake_ws.inject(self.serializer.serialize(msg)) + + receive_task = asyncio.create_task(task._receive_loop()) + await asyncio.sleep(0.1) + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + self.assertEqual(len(sent_to_bus), 1) + self.assertEqual(sent_to_bus[0].target, "worker") + + async def test_drops_inbound_for_other_tasks(self): + """Inbound messages targeted at other tasks are dropped.""" + fake_ws = FakeStarletteWebSocket() + task = await self._create_server(fake_ws) + + sent_to_bus = [] + original_send = task.send_message + + async def capture_send(message): + sent_to_bus.append(message) + await original_send(message) + + task.send_message = capture_send + + msg = BusDataMessage(source="voice", target="other_task") + fake_ws.inject(self.serializer.serialize(msg)) + + receive_task = asyncio.create_task(task._receive_loop()) + await asyncio.sleep(0.1) + receive_task.cancel() + try: + await receive_task + except asyncio.CancelledError: + pass + + self.assertEqual(len(sent_to_bus), 0) + + +if __name__ == "__main__": + unittest.main()