my_voice_assistant_v2/tests/test_realtime.py

import array
import asyncio

from fastapi.testclient import TestClient

import app.dependencies as deps
from app.main import app
from app.audio.vad import rms, EnergyVAD

client = TestClient(app)


def _pcm(amplitude: int, n_samples: int) -> bytes:
    return array.array("h", [amplitude] * n_samples).tobytes()


# --- VAD-Unit-Tests --------------------------------------------------------

def test_rms_silence_vs_loud():
    assert rms(b"") == 0.0
    assert rms(_pcm(0, 100)) == 0.0
    assert rms(_pcm(3000, 100)) > 2000


def test_energy_vad_ends_after_speech_then_silence():
    vad = EnergyVAD(sample_rate=16000, threshold=500, silence_ms=300)
    loud = _pcm(3000, 1600)   # 100 ms Sprache
    silent = _pcm(0, 1600)    # 100 ms Stille
    assert vad.feed(loud) is False
    assert vad.feed(silent) is False   # 100 ms
    assert vad.feed(silent) is False   # 200 ms
    assert vad.feed(silent) is True    # 300 ms -> Ende


def test_energy_vad_ignores_silence_without_speech():
    vad = EnergyVAD(sample_rate=16000, threshold=500, silence_ms=100)
    for _ in range(10):
        assert vad.feed(_pcm(0, 1600)) is False


# --- Barge-in --------------------------------------------------------------

def _install_slow_stream(monkeypatch):
    class SlowLLM:
        async def complete(self, text, history=None, session_id=None):
            return "fertig"

        async def stream(self, text, history=None, session_id=None):
            for i in range(200):
                await asyncio.sleep(0.005)
                yield f"t{i} "

    class StubTTS:
        async def synthesize(self, text, voice=None, audio_format="pcm"):
            return b"A"

    monkeypatch.setitem(deps.LLM_REGISTRY, "slow", lambda s: SlowLLM())
    monkeypatch.setitem(deps.TTS_REGISTRY, "stub", lambda s: StubTTS())
    return {
        "llm_provider": "slow",
        "tts_provider": "stub",
        "output_endpoint": "loopback",
        "stream": True,
    }


def test_ws_chat_interrupt_cancels_response(monkeypatch):
    base = _install_slow_stream(monkeypatch)
    with client.websocket_connect("/ws/chat") as ws:
        ws.send_json({"text": "Hallo", **base})
        assert ws.receive_json()["type"] == "ack"
        assert ws.receive_json()["type"] == "token"  # Antwort laeuft

        ws.send_json({"type": "interrupt"})

        event = None
        for _ in range(500):
            event = ws.receive_json()
            if event["type"] in ("interrupted", "done"):
                break
        assert event["type"] == "interrupted"  # abgebrochen, nicht fertig


# --- VAD im WebSocket ------------------------------------------------------

def _install_voice_stubs(monkeypatch):
    class STT:
        async def transcribe(self, audio_bytes, fmt, language=None):
            return "ok"

    class LLM:
        async def complete(self, text, history=None, session_id=None):
            return "antwort"

    class TTS:
        async def synthesize(self, text, voice=None, audio_format="pcm"):
            return b"A"

    monkeypatch.setitem(deps.STT_REGISTRY, "s", lambda x: STT())
    monkeypatch.setitem(deps.LLM_REGISTRY, "l", lambda x: LLM())
    monkeypatch.setitem(deps.TTS_REGISTRY, "t", lambda x: TTS())
    return {"stt_provider": "s", "llm_provider": "l", "tts_provider": "t", "output_endpoint": "loopback"}


def test_ws_voice_vad_auto_segments_utterance(monkeypatch):
    opts = _install_voice_stubs(monkeypatch)
    start = {
        "type": "start",
        "vad": True,
        "sample_rate": 16000,
        "vad_silence_ms": 200,
        "format": "pcm",
        **opts,
    }
    loud = _pcm(3000, 1600)    # 100 ms Sprache
    silent = _pcm(0, 1600)     # je 100 ms Stille

    with client.websocket_connect("/ws/voice") as ws:
        ws.send_json(start)
        ws.send_bytes(loud)
        ws.send_bytes(silent)   # 100 ms
        ws.send_bytes(silent)   # 200 ms -> VAD-Ende, Turn startet automatisch

        transcript = ws.receive_json()
        assert transcript["type"] == "transcript" and transcript["text"] == "ok"
        assert ws.receive_json()["type"] == "ack"
        assert ws.receive_json()["type"] == "semantic"
        assert ws.receive_bytes() == b"A"
        assert ws.receive_json()["type"] == "done"