diff --git a/CHANGELOG.md b/CHANGELOG.md index a00c6f45b..97edb7382 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -110,6 +110,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Other +- Improved Unit Test `run_test()` to use `PipelineTask` and + `PipelineRunner`. There's now also some control around `StartFrame` and + `EndFrame`. The `EndTaskFrame` has been removed since it doesn't seem + necessary with this new approach. + - Updated `twilio-chatbot` with a few new features: use 8000 sample rate and avoid resampling, a new client useful for stress testing and testing locally without the need to make phone calls. Also, added audio recording on both the diff --git a/src/pipecat/tests/utils.py b/src/pipecat/tests/utils.py index 62f1c7571..ccfca9304 100644 --- a/src/pipecat/tests/utils.py +++ b/src/pipecat/tests/utils.py @@ -5,24 +5,25 @@ # import asyncio -from dataclasses import dataclass -from typing import Awaitable, Callable, Sequence, Tuple +import sys +from typing import Any, Awaitable, Callable, Dict, Sequence, Tuple + +from loguru import logger -from pipecat.clocks.system_clock import SystemClock from pipecat.frames.frames import ( - ControlFrame, + EndFrame, Frame, HeartbeatFrame, StartFrame, ) from pipecat.observers.base_observer import BaseObserver +from pipecat.pipeline.pipeline import Pipeline +from pipecat.pipeline.runner import PipelineRunner +from pipecat.pipeline.task import PipelineParams, PipelineTask from pipecat.processors.frame_processor import FrameDirection, FrameProcessor -from pipecat.utils.asyncio import TaskManager - -@dataclass -class EndTestFrame(ControlFrame): - pass +logger.remove(0) +logger.add(sys.stderr, level="TRACE") class HeartbeatsObserver(BaseObserver): @@ -48,54 +49,58 @@ class HeartbeatsObserver(BaseObserver): class QueuedFrameProcessor(FrameProcessor): - def __init__(self, queue: asyncio.Queue, ignore_start: bool = True): + def __init__( + self, queue: asyncio.Queue, queue_direction: FrameDirection, ignore_start: bool = True + ): super().__init__() self._queue = queue + self._queue_direction = queue_direction self._ignore_start = ignore_start async def process_frame(self, frame: Frame, direction: FrameDirection): await super().process_frame(frame, direction) - if self._ignore_start and isinstance(frame, StartFrame): - await self.push_frame(frame, direction) - else: - await self._queue.put(frame) - await self.push_frame(frame, direction) + if direction == self._queue_direction: + if not isinstance(frame, StartFrame) or not self._ignore_start: + await self._queue.put(frame) + await self.push_frame(frame, direction) async def run_test( processor: FrameProcessor, + *, frames_to_send: Sequence[Frame], expected_down_frames: Sequence[type], expected_up_frames: Sequence[type] = [], + ignore_start: bool = True, + start_metadata: Dict[str, Any] = {}, + send_end_frame: bool = True, ) -> Tuple[Sequence[Frame], Sequence[Frame]]: received_up = asyncio.Queue() received_down = asyncio.Queue() - source = QueuedFrameProcessor(received_up) - sink = QueuedFrameProcessor(received_down) + source = QueuedFrameProcessor(received_up, FrameDirection.UPSTREAM, ignore_start) + sink = QueuedFrameProcessor(received_down, FrameDirection.DOWNSTREAM, ignore_start) - source.link(processor) - processor.link(sink) + pipeline = Pipeline([source, processor, sink]) - task_manager = TaskManager() - task_manager.set_event_loop(asyncio.get_event_loop()) - await source.queue_frame(StartFrame(clock=SystemClock(), task_manager=task_manager)) + task = PipelineTask(pipeline, params=PipelineParams(start_metadata=start_metadata)) for frame in frames_to_send: - await processor.process_frame(frame, FrameDirection.DOWNSTREAM) + await task.queue_frame(frame) - await processor.queue_frame(EndTestFrame()) - await processor.queue_frame(EndTestFrame(), FrameDirection.UPSTREAM) + if send_end_frame: + await task.queue_frame(EndFrame()) + + runner = PipelineRunner() + await runner.run(task) # # Down frames # received_down_frames: Sequence[Frame] = [] - running = True - while running: + while not received_down.empty(): frame = await received_down.get() - running = not isinstance(frame, EndTestFrame) - if running: + if not isinstance(frame, EndFrame) or not send_end_frame: received_down_frames.append(frame) print("received DOWN frames =", received_down_frames) @@ -109,12 +114,9 @@ async def run_test( # Up frames # received_up_frames: Sequence[Frame] = [] - running = True - while running: + while not received_up.empty(): frame = await received_up.get() - running = not isinstance(frame, EndTestFrame) - if running: - received_up_frames.append(frame) + received_up_frames.append(frame) print("received UP frames =", received_up_frames) diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py index 5eb6cc40f..ab7210e49 100644 --- a/tests/test_aggregators.py +++ b/tests/test_aggregators.py @@ -31,7 +31,11 @@ class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase): expected_returned_frames = [TextFrame, TextFrame, TextFrame] - (received_down, _) = await run_test(aggregator, frames_to_send, expected_returned_frames) + (received_down, _) = await run_test( + aggregator, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) assert received_down[-3].text == "Hello, world. " assert received_down[-2].text == "How are you? " assert received_down[-1].text == "I am fine! " @@ -66,5 +70,7 @@ class TestGatedAggregator(unittest.IsolatedAsyncioTestCase): ] (received_down, _) = await run_test( - gated_aggregator, frames_to_send, expected_returned_frames + gated_aggregator, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, ) diff --git a/tests/test_filters.py b/tests/test_filters.py index 4a0fa6a25..9d3b0f003 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -4,7 +4,6 @@ # SPDX-License-Identifier: BSD 2-Clause License # -import asyncio import unittest from pipecat.frames.frames import ( @@ -19,7 +18,7 @@ from pipecat.processors.filters.frame_filter import FrameFilter from pipecat.processors.filters.function_filter import FunctionFilter from pipecat.processors.filters.identity_filter import IdentityFilter from pipecat.processors.filters.wake_check_filter import WakeCheckFilter -from pipecat.tests.utils import EndTestFrame, run_test +from pipecat.tests.utils import run_test class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase): @@ -27,27 +26,44 @@ class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase): filter = IdentityFilter() frames_to_send = [UserStartedSpeakingFrame(), UserStoppedSpeakingFrame()] expected_returned_frames = [UserStartedSpeakingFrame, UserStoppedSpeakingFrame] - await run_test(filter, frames_to_send, expected_returned_frames) + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) class TestFrameFilter(unittest.IsolatedAsyncioTestCase): async def test_text_frame(self): - filter = FrameFilter(types=(TextFrame, EndTestFrame)) + filter = FrameFilter(types=(TextFrame,)) frames_to_send = [TextFrame(text="Hello Pipecat!")] expected_returned_frames = [TextFrame] - await run_test(filter, frames_to_send, expected_returned_frames) + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) async def test_end_frame(self): - filter = FrameFilter(types=(EndFrame, EndTestFrame)) + filter = FrameFilter(types=(EndFrame,)) frames_to_send = [EndFrame()] expected_returned_frames = [EndFrame] - await run_test(filter, frames_to_send, expected_returned_frames) + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + send_end_frame=False, + ) async def test_system_frame(self): - filter = FrameFilter(types=(EndTestFrame,)) + filter = FrameFilter(types=()) frames_to_send = [UserStartedSpeakingFrame()] expected_returned_frames = [UserStartedSpeakingFrame] - await run_test(filter, frames_to_send, expected_returned_frames) + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) class TestFunctionFilter(unittest.IsolatedAsyncioTestCase): @@ -58,7 +74,11 @@ class TestFunctionFilter(unittest.IsolatedAsyncioTestCase): filter = FunctionFilter(filter=passthrough) frames_to_send = [TextFrame(text="Hello Pipecat!")] expected_returned_frames = [TextFrame] - await run_test(filter, frames_to_send, expected_returned_frames) + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) async def test_no_passthrough(self): async def no_passthrough(frame: Frame): @@ -66,14 +86,12 @@ class TestFunctionFilter(unittest.IsolatedAsyncioTestCase): filter = FunctionFilter(filter=no_passthrough) frames_to_send = [TextFrame(text="Hello Pipecat!")] - expected_returned_frames = [TextFrame] - try: - await asyncio.wait_for( - run_test(filter, frames_to_send, expected_returned_frames), timeout=0.5 - ) - assert False - except asyncio.TimeoutError: - pass + expected_returned_frames = [] + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase): @@ -81,7 +99,11 @@ class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase): filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) frames_to_send = [TranscriptionFrame(user_id="test", text="Phrase 1", timestamp="")] expected_returned_frames = [] - await run_test(filter, frames_to_send, expected_returned_frames) + await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) async def test_wake_word(self): filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) @@ -90,5 +112,9 @@ class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase): TranscriptionFrame(user_id="test", text="Phrase 1", timestamp=""), ] expected_returned_frames = [TranscriptionFrame, TranscriptionFrame] - (received_down, _) = await run_test(filter, frames_to_send, expected_returned_frames) + (received_down, _) = await run_test( + filter, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) assert received_down[-1].text == "Phrase 1" diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 6c001b7e0..caac89c11 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -7,7 +7,7 @@ import asyncio import unittest -from pipecat.frames.frames import EndFrame, HeartbeatFrame, TextFrame +from pipecat.frames.frames import EndFrame, HeartbeatFrame, StartFrame, TextFrame from pipecat.pipeline.parallel_pipeline import ParallelPipeline from pipecat.pipeline.pipeline import Pipeline from pipecat.pipeline.task import PipelineParams, PipelineTask @@ -22,7 +22,11 @@ class TestPipeline(unittest.IsolatedAsyncioTestCase): frames_to_send = [TextFrame(text="Hello from Pipecat!")] expected_returned_frames = [TextFrame] - await run_test(pipeline, frames_to_send, expected_returned_frames) + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) async def test_pipeline_multiple(self): identity1 = IdentityFilter() @@ -33,7 +37,25 @@ class TestPipeline(unittest.IsolatedAsyncioTestCase): frames_to_send = [TextFrame(text="Hello from Pipecat!")] expected_returned_frames = [TextFrame] - await run_test(pipeline, frames_to_send, expected_returned_frames) + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) + + async def test_pipeline_start_metadata(self): + pipeline = Pipeline([IdentityFilter()]) + + frames_to_send = [] + expected_returned_frames = [StartFrame] + (received_down, _) = await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ignore_start=False, + start_metadata={"foo": "bar"}, + ) + assert "foo" in received_down[-1].metadata class TestParallelPipeline(unittest.IsolatedAsyncioTestCase): @@ -42,7 +64,11 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase): frames_to_send = [TextFrame(text="Hello from Pipecat!")] expected_returned_frames = [TextFrame] - await run_test(pipeline, frames_to_send, expected_returned_frames) + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) async def test_parallel_multiple(self): """Should only passthrough one instance of TextFrame.""" @@ -50,7 +76,11 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase): frames_to_send = [TextFrame(text="Hello from Pipecat!")] expected_returned_frames = [TextFrame] - await run_test(pipeline, frames_to_send, expected_returned_frames) + await run_test( + pipeline, + frames_to_send=frames_to_send, + expected_down_frames=expected_returned_frames, + ) class TestPipelineTask(unittest.IsolatedAsyncioTestCase): @@ -79,7 +109,9 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase): task = PipelineTask( pipeline, params=PipelineParams( - enable_heartbeats=True, heartbeats_period_secs=0.2, observers=[heartbeats_observer] + enable_heartbeats=True, + heartbeats_period_secs=0.2, + observers=[heartbeats_observer], ), ) task.set_event_loop(asyncio.get_event_loop())