Merge pull request #1048 from pipecat-ai/aleix/add-unittest-utils

tests: add some initial run_test() utilities

src/pipecat/pipeline/base_task.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2024–2025, Daily
+#
+# SPDX-License-Identifier: BSD 2-Clause License
+#
+
+from abc import ABC, abstractmethod
+from typing import AsyncIterable, Iterable
+
+from pipecat.frames.frames import Frame
+
+
+class BaseTask(ABC):
+    @abstractmethod
+    def has_finished(self) -> bool:
+        """Indicates whether the tasks has finished. That is, all processors
+        have stopped.
+
+        """
+        pass
+
+    @abstractmethod
+    async def stop_when_done(self):
+        """This is a helper function that sends an EndFrame to the pipeline in
+        order to stop the task after everything in it has been processed.
+
+        """
+        pass
+
+    @abstractmethod
+    async def cancel(self):
+        """
+        Stops the running pipeline immediately.
+        """
+        pass
+
+    @abstractmethod
+    async def run(self):
+        """
+        Starts running the given pipeline.
+        """
+        pass
+
+    @abstractmethod
+    async def queue_frame(self, frame: Frame):
+        """
+        Queue a frame to be pushed down the pipeline.
+        """
+        pass
+
+    @abstractmethod
+    async def queue_frames(self, frames: Iterable[Frame] | AsyncIterable[Frame]):
+        """
+        Queues multiple frames to be pushed down the pipeline.
+        """
+        pass

src/pipecat/pipeline/task.py

@@ -27,6 +27,7 @@ from pipecat.frames.frames import (
 from pipecat.metrics.metrics import ProcessingMetricsData, TTFBMetricsData
 from pipecat.observers.base_observer import BaseObserver
 from pipecat.pipeline.base_pipeline import BasePipeline
+from pipecat.pipeline.base_task import BaseTask
 from pipecat.pipeline.task_observer import TaskObserver
 from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
 from pipecat.utils.utils import obj_count, obj_id
@@ -45,6 +46,7 @@ class PipelineParams(BaseModel):
     send_initial_empty_metrics: bool = True
     report_only_initial_ttfb: bool = False
     observers: List[BaseObserver] = []
+    heartbeats_period_secs: float = HEARTBEAT_SECONDS
 
 
 class Source(FrameProcessor):
@@ -85,7 +87,7 @@ class Sink(FrameProcessor):
         await self._down_queue.put(frame)
 
 
-class PipelineTask:
+class PipelineTask(BaseTask):
     def __init__(
         self,
         pipeline: BasePipeline,
@@ -121,7 +123,7 @@ class PipelineTask:
 
         self._observer = TaskObserver(params.observers)
 
-    def has_finished(self):
+    def has_finished(self) -> bool:
         """Indicates whether the tasks has finished. That is, all processors
         have stopped.
 
@@ -315,7 +317,7 @@ class PipelineTask:
 
     async def _heartbeat_push_handler(self):
         """
-        This tasks pushes a heartbeat frame every HEARTBEAT_SECONDS.
+        This tasks pushes a heartbeat frame every heartbeat period.
         """
         while True:
             try:
@@ -323,7 +325,7 @@ class PipelineTask:
                 # task will just stop waiting for the pipeline to finish not
                 # allowing more frames to be pushed.
                 await self._source.queue_frame(HeartbeatFrame(timestamp=self._clock.get_time()))
-                await asyncio.sleep(HEARTBEAT_SECONDS)
+                await asyncio.sleep(self._params.heartbeats_period_secs)
             except asyncio.CancelledError:
                 break
 

src/pipecat/processors/filters/frame_filter.py

@@ -11,7 +11,7 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
 
 
 class FrameFilter(FrameProcessor):
-    def __init__(self, types: Tuple[Type[Frame]]):
+    def __init__(self, types: Tuple[Type[Frame], ...]):
         super().__init__()
         self._types = types
 

src/pipecat/processors/frame_processor.py

@@ -293,8 +293,7 @@ class FrameProcessor:
         await self.__input_frame_task
 
     async def __input_frame_task_handler(self):
-        running = True
-        while running:
+        while True:
             try:
                 if self.__should_block_frames:
                     logger.trace(f"{self}: frame processing paused")
@@ -311,8 +310,6 @@ class FrameProcessor:
                 if callback:
                     await callback(self, frame, direction)
 
-                running = not isinstance(frame, EndFrame)
-
                 self.__input_queue.task_done()
             except asyncio.CancelledError:
                 logger.trace(f"{self}: cancelled input task")
@@ -330,12 +327,10 @@ class FrameProcessor:
         await self.__push_frame_task
 
     async def __push_frame_task_handler(self):
-        running = True
-        while running:
+        while True:
             try:
                 (frame, direction) = await self.__push_queue.get()
                 await self.__internal_push_frame(frame, direction)
-                running = not isinstance(frame, EndFrame)
                 self.__push_queue.task_done()
             except asyncio.CancelledError:
                 logger.trace(f"{self}: cancelled push task")

src/pipecat/serializers/protobuf.py

@@ -93,11 +93,11 @@ class ProtobufFrameSerializer(FrameSerializer):
         id = getattr(args, "id", None)
         name = getattr(args, "name", None)
         pts = getattr(args, "pts", None)
-        if not id and "id" in args_dict:
+        if "id" in args_dict:
             del args_dict["id"]
-        if not name and "name" in args_dict:
+        if "name" in args_dict:
             del args_dict["name"]
-        if not pts and "pts" in args_dict:
+        if "pts" in args_dict:
             del args_dict["pts"]
 
         # Create the instance
@@ -105,10 +105,10 @@ class ProtobufFrameSerializer(FrameSerializer):
 
         # Set special fields
         if id:
-            setattr(instance, "id", getattr(args, "id", None))
+            setattr(instance, "id", id)
         if name:
-            setattr(instance, "name", getattr(args, "name", None))
+            setattr(instance, "name", name)
         if pts:
-            setattr(instance, "pts", getattr(args, "pts", None))
+            setattr(instance, "pts", pts)
 
         return instance