esp32-server/tests/test_app.py

"""AC-002: app imports and the /ws/xiaozhi endpoint is live (handshake + turn)."""
from __future__ import annotations

import json

import pytest
from fastapi.testclient import TestClient

from app.config import Config
from app.main import make_app
from app.security import DeviceAuth

TOKEN = "test-device-token"


def _cfg() -> Config:
    return Config.from_dict(
        {
            "security": {"enabled": True, "devices": [
                {"device_id": "d1", "token_hash": DeviceAuth.hash_token(TOKEN)}
            ]},
        }
    )


def test_health():
    client = TestClient(make_app(_cfg()))
    r = client.get("/health")
    assert r.status_code == 200
    assert r.json()["status"] == "ok"


def test_ota_endpoint_delivers_server():
    """FIRMWARE: device reads websocket.url/token from /ota POST."""
    cfg = _cfg()
    from app.config import OtaConfig

    cfg.ota = OtaConfig(
        server_ws_url="wss://zhi.moreminimore.com/ws/xiaozhi",
        server_token=TOKEN,
    )
    client = TestClient(make_app(cfg))
    r = client.post("/ota", json={"device_id": "d1"})
    assert r.status_code == 200
    body = r.json()
    assert body["websocket"]["url"] == "wss://zhi.moreminimore.com/ws/xiaozhi"
    assert body["websocket"]["token"] == TOKEN
    assert "server_time" in body


def test_ws_auth_via_headers():
    """FIRMWARE: token in Authorization header, Device-Id = MAC."""
    client = TestClient(make_app(_cfg()))
    with client.websocket_connect(
        "/ws/xiaozhi",
        headers={
            "Device-Id": "d1",
            "Authorization": f"Bearer {TOKEN}",
        },
    ) as ws:
        # hello frame WITHOUT auth fields — auth came in headers
        ws.send_text(json.dumps({"device_id": "d1"}))
        reply = json.loads(ws.receive_text())
        assert reply["type"] == "hello"
        assert reply["transport"] == "websocket"
        assert reply["session_id"]


def test_ws_header_auth_rejects_bad_token():
    client = TestClient(make_app(_cfg()))
    with client.websocket_connect(
        "/ws/xiaozhi",
        headers={"Device-Id": "d1", "Authorization": "Bearer wrong"},
    ) as ws:
        ws.send_text(json.dumps({"device_id": "d1"}))
        reply = json.loads(ws.receive_text())
        assert reply["type"] == "error"


def test_long_reply_does_not_deadlock():
    """Critical 1 regression: a reply whose audio far exceeds the old
    barge-in TTSQueue cap (64) must still complete. Previously the queue was
    filled with no consumer, so a long reply deadlocked after 64 chunks and
    ``tts_stop`` was never sent. The reply path no longer feeds the queue.
    """
    import asyncio

    from app.audio.vad import make_vad
    from app.metrics import LatencyLogger
    from app.stt import STTResult
    from app.tts import MockTTS
    from app.xiaozhi.session import Session

    class _LongClient:
        def session(self):
            class _S:
                def push(self, u, a):
                    pass
            return _S()

        async def respond(self, text, *, session_search_fn=None):
            for i in range(80):  # 80 sentences -> 160 chunks > 64
                yield f"ประโยค {i}."

    class _STT:
        async def transcribe(self, audio):
            return STTResult(text="สวัสดีครับ")

    session = Session(
        client=_LongClient(),
        stt=_STT(),
        tts=MockTTS(),
        vad=make_vad("energy"),
        logger=LatencyLogger(),
    )

    async def _drain():
        audio_chunks = 0
        async for kind, payload in session.run_turn(b"\x00" * 640):
            if kind == "audio":
                audio_chunks += 1
        return audio_chunks

    # wait_for guards against the deadlock: if it were back, this raises
    # TimeoutError (and the test fails) instead of hanging forever.
    audio_chunks = asyncio.run(asyncio.wait_for(_drain(), timeout=5))
    assert audio_chunks > 64, (
        f"expected a long reply (>64 chunks), got {audio_chunks}"
    )


def test_encoder_is_rate_generic():
    """Critical 2 regression: the reply encoder must be rate-generic, not
    hardwired to 16 kHz. At 24 kHz (the TTS/codec output rate this board
    uses) a 20 ms frame is 480 samples = 960 B — feeding that must yield
    exactly one Opus packet."""
    from app.audio.opus import OpusEncoder, _ensure_native

    # opuslib raises at import when the native lib is not on PATH, so we
    # must make it discoverable BEFORE importing it (mirrors the bridge).
    _ensure_native()
    try:
        import opuslib  # noqa: F401
    except Exception:
        pytest.skip("native opus not available on this host")

    enc = OpusEncoder(sample_rate=24000, channels=1)
    assert enc.FRAME_SAMPLES == 480
    assert enc.FRAME_BYTES == 960
    out = enc.feed(b"\x00" * 960)
    assert len(out) == 1
    assert isinstance(out[0], (bytes, bytearray))


def test_ws_handshake_and_turn():
    """One spoken utterance (burst of speech frames + trailing silence)
    -> firmware-compatible reply: stt text, tts start/sentence/stop, audio."""
    client = TestClient(make_app(_cfg()))
    with client.websocket_connect("/ws/xiaozhi") as ws:
        # hello frame
        ws.send_text(json.dumps({"device_id": "d1", "authorization": TOKEN}))
        reply = json.loads(ws.receive_text())
        assert reply["type"] == "hello"
        assert reply["session_id"]
        # one utterance: 6 loud frames (16-bit, RMS >= 300), then silence
        # frames that push the VAD past min_silence_frames.
        loud = b"\x10\x0f" * 320  # 20 ms @ 16 kHz mono, amplitude 3856
        silent = b"\x00\x00" * 320
        for _ in range(6):
            ws.send_bytes(loud)
        for _ in range(12):
            ws.send_bytes(silent)
        got_stt = got_audio = False
        tts_stop = False
        deadline = 0
        while not tts_stop and deadline < 256:
            frame = ws.receive()
            deadline += 1
            if frame.get("text") is not None:
                payload = json.loads(frame["text"])
                if payload.get("type") == "stt":
                    got_stt = True
                if payload.get("type") == "tts" and payload.get("state") == "stop":
                    tts_stop = True
            else:
                got_audio = True
        assert got_stt and got_audio and tts_stop