137 lines
4.3 KiB
Python
137 lines
4.3 KiB
Python
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 starlette.requests import Request
|
|
|
|
|
|
class FakeConnection:
|
|
pc_id = "pc_test"
|
|
|
|
|
|
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)
|
|
|
|
|
|
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)
|
|
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_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()
|