diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py index 463c2ffed..cc03fda52 100644 --- a/tests/test_aggregators.py +++ b/tests/test_aggregators.py @@ -4,125 +4,67 @@ # SPDX-License-Identifier: BSD 2-Clause License # -import asyncio -import doctest -import functools import unittest from pipecat.frames.frames import ( - AudioRawFrame, - EndFrame, - Frame, ImageRawFrame, LLMFullResponseEndFrame, LLMFullResponseStartFrame, + OutputAudioRawFrame, + OutputImageRawFrame, TextFrame, ) -from pipecat.pipeline.parallel_pipeline import ParallelPipeline -from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.aggregators.gated import GatedAggregator from pipecat.processors.aggregators.sentence import SentenceAggregator -from pipecat.processors.text_transformer import StatelessTextTransformer +from tests.utils import run_test -class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase): - @unittest.skip("FIXME: This test is failing") +class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase): async def test_sentence_aggregator(self): - sentence = "Hello, world. How are you? I am fine" - expected_sentences = ["Hello, world.", " How are you?", " I am fine "] aggregator = SentenceAggregator() + + sentence = "Hello, world. How are you? I am fine!" + + frames_to_send = [] for word in sentence.split(" "): - async for sentence in aggregator.process_frame(TextFrame(word + " ")): - self.assertIsInstance(sentence, TextFrame) - if isinstance(sentence, TextFrame): - self.assertEqual(sentence.text, expected_sentences.pop(0)) + frames_to_send.append(TextFrame(text=word + " ")) - async for sentence in aggregator.process_frame(EndFrame()): - if len(expected_sentences): - self.assertIsInstance(sentence, TextFrame) - if isinstance(sentence, TextFrame): - self.assertEqual(sentence.text, expected_sentences.pop(0)) - else: - self.assertIsInstance(sentence, EndFrame) + expected_returned_frames = [TextFrame, TextFrame, TextFrame] - self.assertEqual(expected_sentences, []) + (received_down, _) = await run_test(aggregator, frames_to_send, 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! " - @unittest.skip("FIXME: This test is failing") - async def test_gated_accumulator(self): + +class TestGatedAggregator(unittest.IsolatedAsyncioTestCase): + async def test_gated_aggregator(self): gated_aggregator = GatedAggregator( gate_open_fn=lambda frame: isinstance(frame, ImageRawFrame), gate_close_fn=lambda frame: isinstance(frame, LLMFullResponseStartFrame), start_open=False, ) - frames = [ + frames_to_send = [ LLMFullResponseStartFrame(), TextFrame("Hello, "), TextFrame("world."), - AudioRawFrame(b"hello"), - ImageRawFrame(b"image", (0, 0)), - AudioRawFrame(b"world"), + OutputAudioRawFrame(audio=b"hello", sample_rate=16000, num_channels=1), + OutputImageRawFrame(image=b"image", size=(0, 0), format="RGB"), + OutputAudioRawFrame(audio=b"world", sample_rate=16000, num_channels=1), LLMFullResponseEndFrame(), ] - expected_output_frames = [ - ImageRawFrame(b"image", (0, 0)), - LLMFullResponseStartFrame(), - TextFrame("Hello, "), - TextFrame("world."), - AudioRawFrame(b"hello"), - AudioRawFrame(b"world"), - LLMFullResponseEndFrame(), + expected_returned_frames = [ + OutputImageRawFrame, + LLMFullResponseStartFrame, + TextFrame, + TextFrame, + OutputAudioRawFrame, + OutputAudioRawFrame, + LLMFullResponseEndFrame, ] - for frame in frames: - async for out_frame in gated_aggregator.process_frame(frame): - self.assertEqual(out_frame, expected_output_frames.pop(0)) - self.assertEqual(expected_output_frames, []) - @unittest.skip("FIXME: This test is failing") - async def test_parallel_pipeline(self): - async def slow_add(sleep_time: float, name: str, x: str): - await asyncio.sleep(sleep_time) - return ":".join([x, name]) - - pipe1_annotation = StatelessTextTransformer(functools.partial(slow_add, 0.1, "pipe1")) - pipe2_annotation = StatelessTextTransformer(functools.partial(slow_add, 0.2, "pipe2")) - sentence_aggregator = SentenceAggregator() - add_dots = StatelessTextTransformer(lambda x: x + ".") - - source = asyncio.Queue() - sink = asyncio.Queue() - pipeline = Pipeline( - [ - ParallelPipeline([[pipe1_annotation], [sentence_aggregator, pipe2_annotation]]), - add_dots, - ], - source, - sink, + (received_down, _) = await run_test( + gated_aggregator, frames_to_send, expected_returned_frames ) - - frames = [TextFrame("Hello, "), TextFrame("world."), EndFrame()] - - expected_output_frames: list[Frame] = [ - TextFrame(text="Hello, :pipe1."), - TextFrame(text="world.:pipe1."), - TextFrame(text="Hello, world.:pipe2."), - EndFrame(), - ] - - for frame in frames: - await source.put(frame) - - await pipeline.run_pipeline() - - while not sink.empty(): - frame = await sink.get() - self.assertEqual(frame, expected_output_frames.pop(0)) - - -def load_tests(loader, tests, ignore): - """Run doctests on the aggregators module.""" - from pipecat.processors import aggregators - - tests.addTests(doctest.DocTestSuite(aggregators)) - return tests