ai-video-fullstack/backend/tests/test_pipeline_events.py

from __future__ import annotations

import unittest
from types import SimpleNamespace
from unittest.mock import patch

from pipecat.frames.frames import (
    BotStartedSpeakingFrame,
    BotStoppedSpeakingFrame,
    OutputTransportMessageUrgentFrame,
)
from services.pipecat.pipeline_events import bind_cascade_pipeline_events


class _EventSource:
    def __init__(self):
        self.handlers = {}

    def event_handler(self, name):
        def decorator(handler):
            self.handlers[name] = handler
            return handler

        return decorator


class _Worker:
    def __init__(self):
        self.frames = []
        self.handlers = {}
        self.downstream_types = set()

    def add_reached_downstream_filter(self, types):
        self.downstream_types.update(types)

    def event_handler(self, name):
        def decorator(handler):
            self.handlers[name] = handler
            return handler

        return decorator

    async def queue_frame(self, frame):
        self.frames.append(frame)


class _Brain:
    spec = SimpleNamespace(owns_context=True)

    def __init__(self, worker):
        self.worker = worker
        self.prepared_greeting = ""
        self.greeting_pending = False
        self.greeting_finished = 0

    def prepare_greeting_context(self, greeting, _context):
        self.prepared_greeting = greeting

    async def on_connected(self, *, greeting_pending=False):
        self.greeting_pending = greeting_pending

    async def on_greeting_finished(self):
        self.greeting_finished += 1

    async def on_client_ready(self):
        for content, timestamp in (
            ("Start Edge 过渡语", "2026-07-14T10:00:00.200+00:00"),
            ("Agent 固定进入语", "2026-07-14T10:00:00.300+00:00"),
        ):
            await self.worker.queue_frame(
                OutputTransportMessageUrgentFrame(
                    message={
                        "type": "transcript",
                        "role": "assistant",
                        "content": content,
                        "timestamp": timestamp,
                    }
                )
            )


class PipelineEventTest(unittest.IsolatedAsyncioTestCase):
    async def test_greeting_keeps_playback_timestamp_until_client_ready(self):
        transport = _EventSource()
        text_input = _EventSource()
        user_aggregator = _EventSource()
        assistant_aggregator = _EventSource()
        worker = _Worker()
        brain = _Brain(worker)

        bind_cascade_pipeline_events(
            transport=transport,
            worker=worker,
            brain=brain,
            context=SimpleNamespace(),
            text_input=text_input,
            user_aggregator=user_aggregator,
            assistant_aggregator=assistant_aggregator,
            greeting="助手开场白",
            vision_enabled=False,
            vision_state={"client_id": None},
        )

        greeting_time = "2026-07-14T10:00:00.100+00:00"
        with patch(
            "services.pipecat.pipeline_events.time_now_iso8601",
            return_value=greeting_time,
        ) as clock:
            await transport.handlers["on_client_connected"](transport, object())
            await text_input.handlers["on_client_ready"](text_input)

        transcripts = [
            frame.message
            for frame in worker.frames
            if isinstance(frame, OutputTransportMessageUrgentFrame)
            and frame.message.get("type") == "transcript"
        ]
        ordered = sorted(transcripts, key=lambda message: message["timestamp"])
        self.assertEqual(
            [message["content"] for message in ordered],
            ["助手开场白", "Start Edge 过渡语", "Agent 固定进入语"],
        )
        self.assertEqual(transcripts[0]["timestamp"], greeting_time)
        self.assertEqual(brain.prepared_greeting, "助手开场白")
        self.assertTrue(brain.greeting_pending)
        clock.assert_called_once_with()

    async def test_greeting_releases_workflow_only_after_real_playback_stop(self):
        transport = _EventSource()
        text_input = _EventSource()
        user_aggregator = _EventSource()
        assistant_aggregator = _EventSource()
        worker = _Worker()
        brain = _Brain(worker)

        bind_cascade_pipeline_events(
            transport=transport,
            worker=worker,
            brain=brain,
            context=SimpleNamespace(),
            text_input=text_input,
            user_aggregator=user_aggregator,
            assistant_aggregator=assistant_aggregator,
            greeting="助手开场白",
            vision_enabled=False,
            vision_state={"client_id": None},
        )

        await transport.handlers["on_client_connected"](transport, object())
        playback_handler = worker.handlers["on_frame_reached_downstream"]

        # An unrelated stop cannot release startup until this greeting has
        # actually produced audio.
        await playback_handler(worker, BotStoppedSpeakingFrame())
        self.assertEqual(brain.greeting_finished, 0)

        await playback_handler(worker, BotStartedSpeakingFrame())
        await playback_handler(worker, BotStoppedSpeakingFrame())
        self.assertEqual(brain.greeting_finished, 1)

        # Duplicate transport notifications are harmless.
        await playback_handler(worker, BotStoppedSpeakingFrame())
        self.assertEqual(brain.greeting_finished, 1)


if __name__ == "__main__":
    unittest.main()