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

from __future__ import annotations

import asyncio
import base64
import json
import unittest
from unittest.mock import AsyncMock, patch

from fastapi import HTTPException
from routes import voice_webrtc
from services.realtime import lifecycle
from starlette.requests import Request


class FakeConnection:
    pc_id = "pc_test"

    def __init__(self):
        self.disconnect = AsyncMock()


def encoded_header_json(value: object) -> str:
    payload = json.dumps(value).encode("utf-8")
    return base64.urlsafe_b64encode(payload).decode("ascii").rstrip("=")


def offer_request(*, client_tools: object | None = None) -> Request:
    headers = []
    if client_tools is not None:
        headers.append(
            (
                b"x-pipecat-client-tools",
                encoded_header_json(client_tools).encode("ascii"),
            )
        )
    return Request({"type": "http", "method": "POST", "headers": headers})


class WebRTCOfferMetadataTest(unittest.IsolatedAsyncioTestCase):
    async def test_client_tool_header_survives_request_endpoint_body_replacement(self):
        request = offer_request(
            client_tools=["set_photo_button_visible", "show_message"]
        )
        payload = {"sdp": "offer", "type": "offer"}

        with patch.object(
            voice_webrtc,
            "_handle_offer_payload",
            new=AsyncMock(return_value={"type": "answer"}),
        ) as handle_offer:
            await voice_webrtc.http_offer(request, payload)

        forwarded = handle_offer.await_args.args[0]
        self.assertEqual(
            forwarded["client_tool_names"],
            ["set_photo_button_visible", "show_message"],
        )

    async def test_client_tool_request_data_is_not_a_capability_channel(self):
        request = offer_request()
        payload = {
            "sdp": "offer",
            "type": "offer",
            "requestData": {"client_tool_names": ["show_message"]},
        }

        with patch.object(
            voice_webrtc,
            "_handle_offer_payload",
            new=AsyncMock(return_value={"type": "answer"}),
        ) as handle_offer:
            await voice_webrtc.http_offer(request, payload)

        forwarded = handle_offer.await_args.args[0]
        self.assertEqual(forwarded["client_tool_names"], [])

    async def test_invalid_client_tool_header_is_rejected(self):
        request = offer_request(client_tools={"name": "show_message"})

        with self.assertRaises(HTTPException) as raised:
            await voice_webrtc.http_offer(
                request,
                {"sdp": "offer", "type": "offer"},
            )

        self.assertEqual(raised.exception.status_code, 400)

    async def test_invalid_assistant_returns_actionable_bad_request(self):
        with patch.object(
            voice_webrtc, "_handle_offer_payload",
            new=AsyncMock(side_effect=ValueError("助手不存在: missing")),
        ):
            with self.assertRaises(HTTPException) as raised:
                await voice_webrtc.http_offer(offer_request(), {"type": "offer"})
        self.assertEqual(raised.exception.status_code, 400)
        self.assertIn("助手不存在", raised.exception.detail)


class WebRTCPipelineLifecycleTest(unittest.IsolatedAsyncioTestCase):
    async def asyncTearDown(self):
        tasks = list(voice_webrtc._pipeline_tasks)
        for task in tasks:
            task.cancel()
        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)
        await asyncio.sleep(0)
        if lifecycle.connection_cleanup_tasks:
            await asyncio.gather(*list(lifecycle.connection_cleanup_tasks))
        voice_webrtc._pipeline_tasks.clear()
        voice_webrtc._connection_tasks.clear()
        voice_webrtc._active_connections.clear()

    async def test_pipeline_task_is_retained_until_completion(self):
        connection = FakeConnection()
        release = asyncio.Event()

        async def pipeline():
            await release.wait()

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        self.assertIn(task, voice_webrtc._pipeline_tasks)
        self.assertIs(voice_webrtc._connection_tasks[connection], task)

        release.set()
        await task
        await asyncio.sleep(0)

        self.assertNotIn(task, voice_webrtc._pipeline_tasks)
        self.assertNotIn(connection, voice_webrtc._connection_tasks)

    async def test_failed_startup_disconnects_and_releases_connection(self):
        connection = FakeConnection()

        async def pipeline():
            raise ValueError("provider construction failed")

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        with self.assertRaises(ValueError):
            await task
        await asyncio.sleep(0)
        if lifecycle.connection_cleanup_tasks:
            await asyncio.gather(*list(lifecycle.connection_cleanup_tasks))
        connection.disconnect.assert_awaited_once()
        self.assertNotIn(connection, voice_webrtc._active_connections)
        self.assertNotIn(connection, voice_webrtc._connection_tasks)

    async def test_cleanup_closed_handler_can_wait_for_completed_pipeline(self):
        connection = FakeConnection()

        async def pipeline():
            return None

        task = voice_webrtc._start_pipeline_task(connection, pipeline())

        async def disconnect():
            await lifecycle.wait_for_pipeline_close(
                task, connection_id=connection.pc_id, timeout=0.01,
            )

        connection.disconnect.side_effect = disconnect
        await task
        await asyncio.sleep(0)
        if lifecycle.connection_cleanup_tasks:
            await asyncio.wait_for(
                asyncio.gather(*list(lifecycle.connection_cleanup_tasks)), 0.5,
            )
        connection.disconnect.assert_awaited_once()

    async def test_already_closed_connection_is_not_disconnected_twice(self):
        connection = FakeConnection()
        release = asyncio.Event()
        task = voice_webrtc._start_pipeline_task(connection, release.wait())
        voice_webrtc._active_connections.discard(connection)
        release.set()
        await task
        await asyncio.sleep(0)
        connection.disconnect.assert_not_awaited()

    async def test_stuck_pipeline_is_cancelled_after_close_grace_period(self):
        connection = FakeConnection()
        cancelled = asyncio.Event()

        async def pipeline():
            try:
                await asyncio.Event().wait()
            finally:
                cancelled.set()

        task = voice_webrtc._start_pipeline_task(connection, pipeline())
        with patch.object(voice_webrtc, "PIPELINE_CLOSE_GRACE_SECONDS", 0.01):
            await voice_webrtc._wait_for_pipeline_close(
                task,
                connection_id=connection.pc_id,
            )

        self.assertTrue(cancelled.is_set())
        self.assertTrue(task.cancelled())


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