tests: add a new SleepFrame
The new SleepFrame allow us to control when system frames are pushed to the pipeline.
This commit is contained in:
@@ -5,6 +5,7 @@
|
|||||||
#
|
#
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
from dataclasses import dataclass
|
||||||
from typing import Any, Awaitable, Callable, Dict, Sequence, Tuple
|
from typing import Any, Awaitable, Callable, Dict, Sequence, Tuple
|
||||||
|
|
||||||
from pipecat.frames.frames import (
|
from pipecat.frames.frames import (
|
||||||
@@ -12,6 +13,7 @@ from pipecat.frames.frames import (
|
|||||||
Frame,
|
Frame,
|
||||||
HeartbeatFrame,
|
HeartbeatFrame,
|
||||||
StartFrame,
|
StartFrame,
|
||||||
|
SystemFrame,
|
||||||
)
|
)
|
||||||
from pipecat.observers.base_observer import BaseObserver
|
from pipecat.observers.base_observer import BaseObserver
|
||||||
from pipecat.pipeline.pipeline import Pipeline
|
from pipecat.pipeline.pipeline import Pipeline
|
||||||
@@ -20,6 +22,15 @@ from pipecat.pipeline.task import PipelineParams, PipelineTask
|
|||||||
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
|
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class SleepFrame(SystemFrame):
|
||||||
|
"""This frame is used by test framework to introduce some sleep time before
|
||||||
|
the next frame is pushed. This is useful to control system frames vs data or
|
||||||
|
control frames."""
|
||||||
|
|
||||||
|
sleep: float = 0.1
|
||||||
|
|
||||||
|
|
||||||
class HeartbeatsObserver(BaseObserver):
|
class HeartbeatsObserver(BaseObserver):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -44,7 +55,11 @@ class HeartbeatsObserver(BaseObserver):
|
|||||||
|
|
||||||
class QueuedFrameProcessor(FrameProcessor):
|
class QueuedFrameProcessor(FrameProcessor):
|
||||||
def __init__(
|
def __init__(
|
||||||
self, queue: asyncio.Queue, queue_direction: FrameDirection, ignore_start: bool = True
|
self,
|
||||||
|
*,
|
||||||
|
queue: asyncio.Queue,
|
||||||
|
queue_direction: FrameDirection,
|
||||||
|
ignore_start: bool = True,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
self._queue = queue
|
self._queue = queue
|
||||||
@@ -72,21 +87,35 @@ async def run_test(
|
|||||||
) -> 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, FrameDirection.UPSTREAM, ignore_start)
|
source = QueuedFrameProcessor(
|
||||||
sink = QueuedFrameProcessor(received_down, FrameDirection.DOWNSTREAM, ignore_start)
|
queue=received_up,
|
||||||
|
queue_direction=FrameDirection.UPSTREAM,
|
||||||
|
ignore_start=ignore_start,
|
||||||
|
)
|
||||||
|
sink = QueuedFrameProcessor(
|
||||||
|
queue=received_down,
|
||||||
|
queue_direction=FrameDirection.DOWNSTREAM,
|
||||||
|
ignore_start=ignore_start,
|
||||||
|
)
|
||||||
|
|
||||||
pipeline = Pipeline([source, processor, sink])
|
pipeline = Pipeline([source, processor, sink])
|
||||||
|
|
||||||
task = PipelineTask(pipeline, params=PipelineParams(start_metadata=start_metadata))
|
task = PipelineTask(pipeline, params=PipelineParams(start_metadata=start_metadata))
|
||||||
|
|
||||||
for frame in frames_to_send:
|
async def push_frames():
|
||||||
await task.queue_frame(frame)
|
# Just give a little head start to the runner.
|
||||||
|
await asyncio.sleep(0.01)
|
||||||
|
for frame in frames_to_send:
|
||||||
|
if isinstance(frame, SleepFrame):
|
||||||
|
await asyncio.sleep(frame.sleep)
|
||||||
|
else:
|
||||||
|
await task.queue_frame(frame)
|
||||||
|
|
||||||
if send_end_frame:
|
if send_end_frame:
|
||||||
await task.queue_frame(EndFrame())
|
await task.queue_frame(EndFrame())
|
||||||
|
|
||||||
runner = PipelineRunner()
|
runner = PipelineRunner()
|
||||||
await runner.run(task)
|
await asyncio.gather(runner.run(task), push_frames())
|
||||||
|
|
||||||
#
|
#
|
||||||
# Down frames
|
# Down frames
|
||||||
|
|||||||
Reference in New Issue
Block a user