102 lines
4.4 KiB
Python
102 lines
4.4 KiB
Python
"""Regression checks for whole-pipeline assembly and terminal outcomes."""
|
|
|
|
import asyncio
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from models import AssistantConfig
|
|
from pipecat.frames.frames import CancelFrame, EndFrame, ErrorFrame, StartFrame
|
|
from pipecat.pipeline.pipeline import Pipeline
|
|
from pipecat.pipeline.worker import PipelineWorker
|
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
|
|
from services.pipecat import pipeline
|
|
from services.pipecat.pipeline_runner import run_media_worker
|
|
|
|
|
|
class PipelineAssemblyTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_realtime_startup_preserves_client_capabilities(self):
|
|
brain = SimpleNamespace(
|
|
spec=SimpleNamespace(supported_runtime_modes={"realtime"}),
|
|
system_prompt=lambda cfg: "test",
|
|
greeting=AsyncMock(return_value=""),
|
|
setup_realtime=AsyncMock(),
|
|
on_session_update=AsyncMock(),
|
|
)
|
|
transport = SimpleNamespace(
|
|
input=lambda: FrameProcessor(),
|
|
output=lambda: FrameProcessor(),
|
|
event_handler=lambda name: lambda handler: handler,
|
|
)
|
|
with (
|
|
patch.object(pipeline, "build_brain", return_value=brain),
|
|
patch.object(pipeline, "create_realtime_service", return_value=FrameProcessor()),
|
|
patch.object(pipeline.ConversationRecorder, "start", new=AsyncMock(return_value=None)),
|
|
patch.object(pipeline, "run_media_worker", new=AsyncMock()) as run,
|
|
):
|
|
await pipeline.run_pipeline(
|
|
transport, AssistantConfig(runtimeMode="realtime"),
|
|
client_tool_names={"show_message"},
|
|
)
|
|
run.assert_awaited_once()
|
|
runtime = brain.setup_realtime.await_args.args[1]
|
|
self.assertTrue(runtime.client_tools.supports("show_message"))
|
|
self.assertFalse(runtime.client_tools.supports("unregistered_tool"))
|
|
|
|
|
|
class TerminalProcessor(FrameProcessor):
|
|
def __init__(self, terminal):
|
|
super().__init__()
|
|
self.terminal = terminal
|
|
|
|
async def process_frame(self, frame, direction):
|
|
await super().process_frame(frame, direction)
|
|
await self.push_frame(frame, direction)
|
|
if isinstance(frame, StartFrame):
|
|
if isinstance(self.terminal, ErrorFrame):
|
|
await self.push_frame(self.terminal, FrameDirection.UPSTREAM)
|
|
else:
|
|
await self.push_frame(self.terminal)
|
|
|
|
|
|
class PipelineOutcomeTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_fatal_error_is_failed_even_when_runner_returns_normally(self):
|
|
worker = PipelineWorker(
|
|
Pipeline([TerminalProcessor(ErrorFrame(error="provider failed", fatal=True))]),
|
|
enable_rtvi=False,
|
|
)
|
|
recorder = SimpleNamespace(finish=AsyncMock())
|
|
await asyncio.wait_for(run_media_worker(worker, recorder), 3)
|
|
recorder.finish.assert_awaited_once_with(status="failed")
|
|
|
|
async def test_normal_end_and_client_disconnect_are_completed(self):
|
|
for terminal in (EndFrame(), CancelFrame(reason="client_disconnected")):
|
|
with self.subTest(terminal=type(terminal).__name__):
|
|
worker = PipelineWorker(Pipeline([Passthrough()]), enable_rtvi=False)
|
|
recorder = SimpleNamespace(finish=AsyncMock())
|
|
await worker.queue_frame(terminal)
|
|
await asyncio.wait_for(run_media_worker(worker, recorder), 3)
|
|
recorder.finish.assert_awaited_once_with(status="completed")
|
|
|
|
async def test_server_task_cancellation_remains_failed_and_propagates(self):
|
|
worker = PipelineWorker(Pipeline([Passthrough()]), enable_rtvi=False)
|
|
started = asyncio.Event()
|
|
|
|
@worker.event_handler("on_pipeline_started")
|
|
async def on_started(_worker, _frame):
|
|
started.set()
|
|
|
|
recorder = SimpleNamespace(finish=AsyncMock())
|
|
task = asyncio.create_task(run_media_worker(worker, recorder))
|
|
await asyncio.wait_for(started.wait(), 3)
|
|
task.cancel()
|
|
with self.assertRaises(asyncio.CancelledError):
|
|
await asyncio.wait_for(task, 3)
|
|
recorder.finish.assert_awaited_once_with(status="failed")
|
|
|
|
|
|
class Passthrough(FrameProcessor):
|
|
async def process_frame(self, frame, direction):
|
|
await super().process_frame(frame, direction)
|
|
await self.push_frame(frame, direction)
|