153 lines
4.3 KiB
Python
153 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import struct
|
|
|
|
from fastapi import WebSocketDisconnect
|
|
|
|
from app.utils.background import BridgeOutputStats, ws_out_loop
|
|
|
|
|
|
class _ClosingWebSocket:
|
|
def __init__(self, exc: BaseException) -> None:
|
|
self.exc = exc
|
|
self.sent_frames = 0
|
|
|
|
async def send_bytes(self, frame: bytes) -> None:
|
|
self.sent_frames += 1
|
|
raise self.exc
|
|
|
|
|
|
class _CollectThenDisconnectWebSocket:
|
|
def __init__(self, disconnect_after: int) -> None:
|
|
self.disconnect_after = disconnect_after
|
|
self.sent_frames: list[bytes] = []
|
|
|
|
async def send_bytes(self, frame: bytes) -> None:
|
|
self.sent_frames.append(frame)
|
|
if len(self.sent_frames) >= self.disconnect_after:
|
|
raise WebSocketDisconnect(code=1000)
|
|
|
|
class _CaptureThenDisconnectWebSocket:
|
|
def __init__(self) -> None:
|
|
self.sent_frames: list[bytes] = []
|
|
|
|
async def send_bytes(self, frame: bytes) -> None:
|
|
self.sent_frames.append(frame)
|
|
raise WebSocketDisconnect(code=1000)
|
|
|
|
|
|
async def _capture_one_agent_frame(frame: bytes, **ws_out_kwargs) -> _CaptureThenDisconnectWebSocket:
|
|
ws = _CaptureThenDisconnectWebSocket()
|
|
agent_q: asyncio.Queue[bytes] = asyncio.Queue()
|
|
await agent_q.put(frame)
|
|
|
|
await asyncio.wait_for(
|
|
ws_out_loop(
|
|
ws,
|
|
agent_q,
|
|
frame_ms=20,
|
|
bytes_per_frame=len(frame),
|
|
**ws_out_kwargs,
|
|
),
|
|
timeout=0.2,
|
|
)
|
|
|
|
return ws
|
|
|
|
|
|
def test_ws_out_loop_exits_when_close_was_already_sent() -> None:
|
|
async def _run() -> None:
|
|
ws = _ClosingWebSocket(RuntimeError('Cannot call "send" once a close message has been sent.'))
|
|
|
|
await asyncio.wait_for(
|
|
ws_out_loop(
|
|
ws,
|
|
asyncio.Queue(),
|
|
frame_ms=20,
|
|
bytes_per_frame=4,
|
|
),
|
|
timeout=0.2,
|
|
)
|
|
|
|
assert ws.sent_frames == 1
|
|
|
|
asyncio.run(_run())
|
|
|
|
|
|
def test_ws_out_loop_applies_default_output_gain(monkeypatch) -> None:
|
|
monkeypatch.delenv("WS_OUTPUT_GAIN", raising=False)
|
|
|
|
frame = struct.pack("<hh", 1000, -1000)
|
|
ws = asyncio.run(_capture_one_agent_frame(frame))
|
|
|
|
assert struct.unpack("<hh", ws.sent_frames[0]) == (1000, -1000)
|
|
|
|
|
|
def test_ws_out_loop_reads_output_gain_from_env(monkeypatch) -> None:
|
|
monkeypatch.setenv("WS_OUTPUT_GAIN", "1.5")
|
|
|
|
frame = struct.pack("<hh", 1000, -1000)
|
|
ws = asyncio.run(_capture_one_agent_frame(frame))
|
|
|
|
assert struct.unpack("<hh", ws.sent_frames[0]) == (1500, -1500)
|
|
|
|
|
|
def test_ws_out_loop_prefers_call_config_output_gain(monkeypatch) -> None:
|
|
monkeypatch.setenv("WS_OUTPUT_GAIN", "1.5")
|
|
|
|
frame = struct.pack("<hh", 1000, -1000)
|
|
ws = asyncio.run(_capture_one_agent_frame(frame, output_gain=1.0))
|
|
|
|
assert struct.unpack("<hh", ws.sent_frames[0]) == (1000, -1000)
|
|
|
|
|
|
def test_ws_out_loop_exits_on_websocket_disconnect() -> None:
|
|
async def _run() -> None:
|
|
ws = _ClosingWebSocket(WebSocketDisconnect(code=1000))
|
|
|
|
await asyncio.wait_for(
|
|
ws_out_loop(
|
|
ws,
|
|
asyncio.Queue(),
|
|
frame_ms=20,
|
|
bytes_per_frame=4,
|
|
),
|
|
timeout=0.2,
|
|
)
|
|
|
|
assert ws.sent_frames == 1
|
|
|
|
asyncio.run(_run())
|
|
|
|
|
|
def test_ws_out_loop_waits_for_speech_before_signaling_first_agent_audio() -> None:
|
|
async def _run() -> None:
|
|
silent_frame = struct.pack("<hh", 0, 0)
|
|
voiced_frame = struct.pack("<hh", 2000, -2000)
|
|
ws = _CollectThenDisconnectWebSocket(disconnect_after=3)
|
|
agent_q: asyncio.Queue[bytes] = asyncio.Queue()
|
|
await agent_q.put(silent_frame)
|
|
await agent_q.put(voiced_frame)
|
|
first_audio = asyncio.Event()
|
|
stats = BridgeOutputStats()
|
|
|
|
await asyncio.wait_for(
|
|
ws_out_loop(
|
|
ws,
|
|
agent_q,
|
|
frame_ms=20,
|
|
bytes_per_frame=len(silent_frame),
|
|
first_agent_audio_sent=first_audio,
|
|
output_stats=stats,
|
|
),
|
|
timeout=0.2,
|
|
)
|
|
|
|
assert ws.sent_frames[:2] == [silent_frame, voiced_frame]
|
|
assert first_audio.is_set()
|
|
assert stats.first_agent_audio_sent is True
|
|
assert stats.agent_audio_bursts == 1
|
|
|
|
asyncio.run(_run())
|