first commit

tests/utils/test_background.py

@@ -0,0 +1,152 @@
+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())