omnivoice-chatterbox-bridge/tests/test_server.py

"""Smoke tests for the OmniVoice <-> Chatterbox bridge."""
from __future__ import annotations

import os
import struct
import unittest
from unittest.mock import patch

from fastapi.testclient import TestClient

from server import create_app


class _Resp:
    def __init__(self, status_code=200, body=b"", json_body=None, content_type="application/json"):
        self.status_code = status_code
        self.content = body
        self._json = json_body
        self.headers = {"Content-Type": content_type}

    def raise_for_status(self):
        if self.status_code >= 400:
            raise RuntimeError(f"http {self.status_code}")

    def json(self):
        return self._json


class _FakeAsyncClient:
    instances: list = []

    def __init__(self, responses):
        self._responses = list(responses)
        self.calls: list[tuple[str, dict]] = []
        _FakeAsyncClient.instances.append(self)

    async def __aenter__(self):
        return self

    async def __aexit__(self, *_):
        return False

    async def get(self, url, headers=None):
        self.calls.append(("GET", {"url": url, "headers": headers}))
        return self._responses.pop(0)

    async def post(self, url, headers=None, json=None):
        self.calls.append(("POST", {"url": url, "headers": headers, "json": json}))
        return self._responses.pop(0)


def _factory(responses):
    def _make(*_args, **_kwargs):
        return _FakeAsyncClient(responses)
    return _make


def _build_wav() -> bytes:
    fmt = struct.pack("<4sIHHIIHH", b"fmt ", 16, 1, 1, 22050, 44100, 2, 16)
    data = struct.pack("<4sI", b"data", 4) + b"\x00\x00\x00\x00"
    return b"RIFF" + struct.pack("<I", 4 + len(fmt) + len(data)) + b"WAVE" + fmt + data


_WAV = _build_wav()


class BridgeTests(unittest.TestCase):
    def setUp(self):
        os.environ["OMNIVOICE_URL"] = "http://upstream.invalid:7860"
        os.environ.pop("OMNIVOICE_API_KEY", None)
        os.environ["OMNIVOICE_VOICE"] = "sparkle"
        _FakeAsyncClient.instances.clear()
        self.client = TestClient(create_app())

    def test_health(self):
        r = self.client.get("/")
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.json()["status"], "ok")
        self.assertEqual(r.json()["provider"], "chatterbox")

    def test_voices_passthrough(self):
        with patch(
            "server.httpx.AsyncClient",
            _factory([_Resp(json_body={
                "sparkle": {"language": "en", "speed": 1.0, "loaded": True},
                "bibi": {"language": "en", "speed": 1.0, "loaded": True},
            })]),
        ):
            r = self.client.get("/api/voices")
        self.assertEqual(r.status_code, 200)
        names = [v["name"] for v in r.json()["voices"]]
        self.assertIn("sparkle", names)
        self.assertIn("bibi", names)

    def test_synthesize_translates_payload_and_returns_wav(self):
        with patch(
            "server.httpx.AsyncClient",
            _factory([_Resp(status_code=200, body=_WAV, content_type="audio/wav")]),
        ):
            r = self.client.post(
                "/v1/audio/speech",
                json={"input": "Hello robot.", "voice": "sparkle"},
            )
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.headers["content-type"], "audio/wav")
        self.assertEqual(r.content[:4], b"RIFF")
        sent = _FakeAsyncClient.instances[0].calls[0]
        self.assertEqual(sent[0], "POST")
        self.assertTrue(sent[1]["url"].endswith("/tts"))
        self.assertEqual(sent[1]["json"]["text"], "Hello robot.")
        self.assertEqual(sent[1]["json"]["voice"], "sparkle")

    def test_synthesize_default_voice_remapped_to_env(self):
        os.environ["OMNIVOICE_VOICE"] = "bibi"
        with patch(
            "server.httpx.AsyncClient",
            _factory([_Resp(status_code=200, body=_WAV, content_type="audio/wav")]),
        ):
            r = self.client.post(
                "/v1/audio/speech",
                json={"input": "Hi.", "voice": "default"},
            )
        self.assertEqual(r.status_code, 200)
        sent_json = _FakeAsyncClient.instances[0].calls[0][1]["json"]
        self.assertEqual(sent_json["voice"], "bibi")

    def test_synthesize_forwards_bearer_token(self):
        os.environ["OMNIVOICE_API_KEY"] = "secret-token-abc"
        with patch(
            "server.httpx.AsyncClient",
            _factory([_Resp(status_code=200, body=_WAV, content_type="audio/wav")]),
        ):
            r = self.client.post(
                "/v1/audio/speech",
                json={"input": "Hi.", "voice": "sparkle"},
            )
        self.assertEqual(r.status_code, 200)
        headers = _FakeAsyncClient.instances[0].calls[0][1]["headers"]
        self.assertEqual(headers.get("Authorization"), "Bearer secret-token-abc")

    def test_synthesize_rejects_empty_text(self):
        r = self.client.post("/v1/audio/speech", json={"input": "  ", "voice": "sparkle"})
        self.assertEqual(r.status_code, 400)

    def test_synthesize_401_from_upstream_becomes_502(self):
        with patch(
            "server.httpx.AsyncClient",
            _factory([_Resp(status_code=401, body=b'{"detail":"bad key"}')]),
        ):
            r = self.client.post(
                "/v1/audio/speech",
                json={"input": "Hi.", "voice": "sparkle"},
            )
        self.assertEqual(r.status_code, 502)

    def test_stream_endpoint_returns_pcm_with_headers(self):
        with patch(
            "server.httpx.AsyncClient",
            _factory([_Resp(status_code=200, body=_WAV, content_type="audio/wav")]),
        ):
            r = self.client.post(
                "/v1/audio/speech/stream",
                json={"input": "stream me", "voice": "sparkle"},
            )
        self.assertEqual(r.status_code, 200)
        self.assertEqual(r.headers["x-encoding"], "pcm_s16le")
        self.assertEqual(r.headers["x-channels"], "1")
        self.assertEqual(r.headers["x-sample-rate"], "22050")


if __name__ == "__main__":
    unittest.main()