tests: run_test() now uses PipelineTask

This commit is contained in:
Aleix Conchillo Flaqué
2025-02-05 10:02:11 -08:00
parent 61f6669926
commit feab9c8fa2
5 changed files with 133 additions and 62 deletions

View File

@@ -110,6 +110,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Other ### 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 - 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 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 without the need to make phone calls. Also, added audio recording on both the

View File

@@ -5,24 +5,25 @@
# #
import asyncio import asyncio
from dataclasses import dataclass import sys
from typing import Awaitable, Callable, Sequence, Tuple 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 ( from pipecat.frames.frames import (
ControlFrame, EndFrame,
Frame, Frame,
HeartbeatFrame, HeartbeatFrame,
StartFrame, StartFrame,
) )
from pipecat.observers.base_observer import BaseObserver 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.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.asyncio import TaskManager
logger.remove(0)
@dataclass logger.add(sys.stderr, level="TRACE")
class EndTestFrame(ControlFrame):
pass
class HeartbeatsObserver(BaseObserver): class HeartbeatsObserver(BaseObserver):
@@ -48,54 +49,58 @@ class HeartbeatsObserver(BaseObserver):
class QueuedFrameProcessor(FrameProcessor): 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__() super().__init__()
self._queue = queue self._queue = queue
self._queue_direction = queue_direction
self._ignore_start = ignore_start self._ignore_start = ignore_start
async def process_frame(self, frame: Frame, direction: FrameDirection): async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction) await super().process_frame(frame, direction)
if self._ignore_start and isinstance(frame, StartFrame): if direction == self._queue_direction:
await self.push_frame(frame, direction) if not isinstance(frame, StartFrame) or not self._ignore_start:
else: await self._queue.put(frame)
await self._queue.put(frame) await self.push_frame(frame, direction)
await self.push_frame(frame, direction)
async def run_test( async def run_test(
processor: FrameProcessor, processor: FrameProcessor,
*,
frames_to_send: Sequence[Frame], frames_to_send: Sequence[Frame],
expected_down_frames: Sequence[type], expected_down_frames: Sequence[type],
expected_up_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]]: ) -> Tuple[Sequence[Frame], Sequence[Frame]]:
received_up = asyncio.Queue() received_up = asyncio.Queue()
received_down = asyncio.Queue() received_down = asyncio.Queue()
source = QueuedFrameProcessor(received_up) source = QueuedFrameProcessor(received_up, FrameDirection.UPSTREAM, ignore_start)
sink = QueuedFrameProcessor(received_down) sink = QueuedFrameProcessor(received_down, FrameDirection.DOWNSTREAM, ignore_start)
source.link(processor) pipeline = Pipeline([source, processor, sink])
processor.link(sink)
task_manager = TaskManager() task = PipelineTask(pipeline, params=PipelineParams(start_metadata=start_metadata))
task_manager.set_event_loop(asyncio.get_event_loop())
await source.queue_frame(StartFrame(clock=SystemClock(), task_manager=task_manager))
for frame in frames_to_send: for frame in frames_to_send:
await processor.process_frame(frame, FrameDirection.DOWNSTREAM) await task.queue_frame(frame)
await processor.queue_frame(EndTestFrame()) if send_end_frame:
await processor.queue_frame(EndTestFrame(), FrameDirection.UPSTREAM) await task.queue_frame(EndFrame())
runner = PipelineRunner()
await runner.run(task)
# #
# Down frames # Down frames
# #
received_down_frames: Sequence[Frame] = [] received_down_frames: Sequence[Frame] = []
running = True while not received_down.empty():
while running:
frame = await received_down.get() frame = await received_down.get()
running = not isinstance(frame, EndTestFrame) if not isinstance(frame, EndFrame) or not send_end_frame:
if running:
received_down_frames.append(frame) received_down_frames.append(frame)
print("received DOWN frames =", received_down_frames) print("received DOWN frames =", received_down_frames)
@@ -109,12 +114,9 @@ async def run_test(
# Up frames # Up frames
# #
received_up_frames: Sequence[Frame] = [] received_up_frames: Sequence[Frame] = []
running = True while not received_up.empty():
while running:
frame = await received_up.get() frame = await received_up.get()
running = not isinstance(frame, EndTestFrame) received_up_frames.append(frame)
if running:
received_up_frames.append(frame)
print("received UP frames =", received_up_frames) print("received UP frames =", received_up_frames)

View File

@@ -31,7 +31,11 @@ class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase):
expected_returned_frames = [TextFrame, TextFrame, TextFrame] 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[-3].text == "Hello, world. "
assert received_down[-2].text == "How are you? " assert received_down[-2].text == "How are you? "
assert received_down[-1].text == "I am fine! " assert received_down[-1].text == "I am fine! "
@@ -66,5 +70,7 @@ class TestGatedAggregator(unittest.IsolatedAsyncioTestCase):
] ]
(received_down, _) = await run_test( (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,
) )

View File

@@ -4,7 +4,6 @@
# SPDX-License-Identifier: BSD 2-Clause License # SPDX-License-Identifier: BSD 2-Clause License
# #
import asyncio
import unittest import unittest
from pipecat.frames.frames import ( 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.function_filter import FunctionFilter
from pipecat.processors.filters.identity_filter import IdentityFilter from pipecat.processors.filters.identity_filter import IdentityFilter
from pipecat.processors.filters.wake_check_filter import WakeCheckFilter 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): class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase):
@@ -27,27 +26,44 @@ class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase):
filter = IdentityFilter() filter = IdentityFilter()
frames_to_send = [UserStartedSpeakingFrame(), UserStoppedSpeakingFrame()] frames_to_send = [UserStartedSpeakingFrame(), UserStoppedSpeakingFrame()]
expected_returned_frames = [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): class TestFrameFilter(unittest.IsolatedAsyncioTestCase):
async def test_text_frame(self): async def test_text_frame(self):
filter = FrameFilter(types=(TextFrame, EndTestFrame)) filter = FrameFilter(types=(TextFrame,))
frames_to_send = [TextFrame(text="Hello Pipecat!")] frames_to_send = [TextFrame(text="Hello Pipecat!")]
expected_returned_frames = [TextFrame] 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): async def test_end_frame(self):
filter = FrameFilter(types=(EndFrame, EndTestFrame)) filter = FrameFilter(types=(EndFrame,))
frames_to_send = [EndFrame()] frames_to_send = [EndFrame()]
expected_returned_frames = [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): async def test_system_frame(self):
filter = FrameFilter(types=(EndTestFrame,)) filter = FrameFilter(types=())
frames_to_send = [UserStartedSpeakingFrame()] frames_to_send = [UserStartedSpeakingFrame()]
expected_returned_frames = [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): class TestFunctionFilter(unittest.IsolatedAsyncioTestCase):
@@ -58,7 +74,11 @@ class TestFunctionFilter(unittest.IsolatedAsyncioTestCase):
filter = FunctionFilter(filter=passthrough) filter = FunctionFilter(filter=passthrough)
frames_to_send = [TextFrame(text="Hello Pipecat!")] frames_to_send = [TextFrame(text="Hello Pipecat!")]
expected_returned_frames = [TextFrame] 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 test_no_passthrough(self):
async def no_passthrough(frame: Frame): async def no_passthrough(frame: Frame):
@@ -66,14 +86,12 @@ class TestFunctionFilter(unittest.IsolatedAsyncioTestCase):
filter = FunctionFilter(filter=no_passthrough) filter = FunctionFilter(filter=no_passthrough)
frames_to_send = [TextFrame(text="Hello Pipecat!")] frames_to_send = [TextFrame(text="Hello Pipecat!")]
expected_returned_frames = [TextFrame] expected_returned_frames = []
try: await run_test(
await asyncio.wait_for( filter,
run_test(filter, frames_to_send, expected_returned_frames), timeout=0.5 frames_to_send=frames_to_send,
) expected_down_frames=expected_returned_frames,
assert False )
except asyncio.TimeoutError:
pass
class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase): class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase):
@@ -81,7 +99,11 @@ class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase):
filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"])
frames_to_send = [TranscriptionFrame(user_id="test", text="Phrase 1", timestamp="")] frames_to_send = [TranscriptionFrame(user_id="test", text="Phrase 1", timestamp="")]
expected_returned_frames = [] 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): async def test_wake_word(self):
filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"])
@@ -90,5 +112,9 @@ class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase):
TranscriptionFrame(user_id="test", text="Phrase 1", timestamp=""), TranscriptionFrame(user_id="test", text="Phrase 1", timestamp=""),
] ]
expected_returned_frames = [TranscriptionFrame, TranscriptionFrame] 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" assert received_down[-1].text == "Phrase 1"

View File

@@ -7,7 +7,7 @@
import asyncio import asyncio
import unittest 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.parallel_pipeline import ParallelPipeline
from pipecat.pipeline.pipeline import Pipeline from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.task import PipelineParams, PipelineTask from pipecat.pipeline.task import PipelineParams, PipelineTask
@@ -22,7 +22,11 @@ class TestPipeline(unittest.IsolatedAsyncioTestCase):
frames_to_send = [TextFrame(text="Hello from Pipecat!")] frames_to_send = [TextFrame(text="Hello from Pipecat!")]
expected_returned_frames = [TextFrame] 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): async def test_pipeline_multiple(self):
identity1 = IdentityFilter() identity1 = IdentityFilter()
@@ -33,7 +37,25 @@ class TestPipeline(unittest.IsolatedAsyncioTestCase):
frames_to_send = [TextFrame(text="Hello from Pipecat!")] frames_to_send = [TextFrame(text="Hello from Pipecat!")]
expected_returned_frames = [TextFrame] 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): class TestParallelPipeline(unittest.IsolatedAsyncioTestCase):
@@ -42,7 +64,11 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase):
frames_to_send = [TextFrame(text="Hello from Pipecat!")] frames_to_send = [TextFrame(text="Hello from Pipecat!")]
expected_returned_frames = [TextFrame] 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): async def test_parallel_multiple(self):
"""Should only passthrough one instance of TextFrame.""" """Should only passthrough one instance of TextFrame."""
@@ -50,7 +76,11 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase):
frames_to_send = [TextFrame(text="Hello from Pipecat!")] frames_to_send = [TextFrame(text="Hello from Pipecat!")]
expected_returned_frames = [TextFrame] 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): class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
@@ -79,7 +109,9 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
task = PipelineTask( task = PipelineTask(
pipeline, pipeline,
params=PipelineParams( 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()) task.set_event_loop(asyncio.get_event_loop())