path: root/app/tests/test_tts_sglomni.py

"""Tests for the SGLang-Omni TTS client (converter/clients/sglomni.py)."""

import base64
import io
import json
import tempfile
import unittest
import urllib.error
import wave
from pathlib import Path
from unittest.mock import MagicMock, patch

from converter import config
from converter.clients import SgOmniTTSClient
from converter.clients.base import NonRetryableTTSError
from converter.clients.sglomni import _data_url, _is_loopback
from converter.clients.speakers import QWEN3_TTS_SPEAKERS


class CatalogConsistencyTests(unittest.TestCase):
    """The backend catalog's vendored facts match the converter's."""

    def test_customvoice_speakers_match_the_qwen_table(self):
        from backends.sglomni.catalog import QWEN_CUSTOMVOICE_SPEAKERS
        self.assertEqual(QWEN_CUSTOMVOICE_SPEAKERS, QWEN3_TTS_SPEAKERS)

_DUMMY_CHUNKS = Path(tempfile.gettempdir()) / "audiobook_sglomni_test_chunks"


def _make_wav() -> bytes:
    """A real minimal RIFF/WAVE file (what a server response looks like)."""
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(24000)
        wav_file.writeframes(b"\x01\x00" * 16)
    return buffer.getvalue()


_WAV_BYTES = _make_wav()
_WAV_FRAMES = b"\x01\x00" * 16


class LoopbackTests(unittest.TestCase):
    def test_loopback_hosts(self):
        self.assertTrue(_is_loopback("http://127.0.0.1:8100"))
        self.assertTrue(_is_loopback("http://localhost:8100"))
        self.assertFalse(_is_loopback("http://10.20.30.40:8100"))

    def test_data_url_carries_mime_and_bytes(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp) / "ref.wav"
            path.write_bytes(b"abc")
            url = _data_url(path)
        self.assertTrue(url.startswith("data:audio/wav;base64,"))
        self.assertEqual(
            base64.b64decode(url.partition(";base64,")[2]), b"abc")


class ConnectInputTests(unittest.TestCase):
    """Capability-driven validation before any HTTP is attempted."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.ref = Path(self._tmp.name) / "narrator.wav"
        self.ref.write_bytes(b"abc")
        self.addCleanup(self._tmp.cleanup)

    def _client(self, model="higgs_audio_v3_tts", **kwargs):
        # Bypass _connect (HTTP) — these tests cover the input checks.
        with patch.object(SgOmniTTSClient, "_connect"):
            return SgOmniTTSClient(_DUMMY_CHUNKS, model=model, **kwargs)

    def test_unknown_model_raises(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="nope")
        self.assertIn("Unknown SGLang-Omni model", str(ctx.exception))

    def test_design_model_requires_instructions(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="qwen3_tts_1_7b_voicedesign")
        self.assertIn("--instructions", str(ctx.exception))

    def test_reference_required_model_refuses_to_connect_without_one(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="qwen3_tts_1_7b_base")
        self.assertIn("requires reference audio", str(ctx.exception))

    def test_missing_reference_file_raises(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._client(model="higgs_audio_v3_tts",
                         ref_audio=str(Path(self._tmp.name) / "gone.wav"))
        self.assertIn("Reference audio not found", str(ctx.exception))

    def test_clone_capable_model_allows_text_only(self):
        client = self._client(model="higgs_audio_v3_tts")
        self.assertIsNone(client.ref_audio)

    def test_speaker_model_ignores_the_clone_reference(self):
        client = self._client(model="qwen3_tts_0_6b_customvoice",
                              ref_audio=str(self.ref))
        self.assertIsNone(client.ref_audio)

    def test_seed_only_sent_for_models_that_accept_it(self):
        with patch("converter.clients.sglomni.resolve_request_seed",
                   return_value=42):
            client = self._client(model="qwen3_tts_1_7b_base",
                                  ref_audio=str(self.ref))
            self.assertEqual(client._seed, 42)
            client = self._client(model="higgs_audio_v3_tts")
            self.assertIsNone(client._seed)

    def test_negative_seed_is_not_sent(self):
        with patch("converter.clients.sglomni.resolve_request_seed",
                   return_value=-1):
            client = self._client(model="qwen3_tts_1_7b_base",
                                  ref_audio=str(self.ref))
            self.assertIsNone(client._seed)


class ConnectHealthTests(unittest.TestCase):
    """_connect gates on /health and the hosted model."""

    def _response(self, payload):
        response = MagicMock()
        response.__enter__.return_value = response
        response.read.return_value = json.dumps(payload).encode("utf-8")
        return response

    def _connect(self, payloads, **kwargs):
        # urlopen is called once per _get_json call, in order.
        responses = [self._response(payload) for payload in payloads]
        with patch("converter.clients.sglomni.urllib.request.urlopen",
                   side_effect=responses):
            with patch.object(SgOmniTTSClient, "_resolve_reference_text"):
                return SgOmniTTSClient(_DUMMY_CHUNKS,
                                       model="higgs_audio_v3_tts", **kwargs)

    def test_unreachable_server_raises_with_guidance(self):
        import urllib.error
        with patch("converter.clients.sglomni.urllib.request.urlopen",
                   side_effect=urllib.error.URLError("refused")):
            with self.assertRaises(RuntimeError) as ctx:
                SgOmniTTSClient(_DUMMY_CHUNKS, model="higgs_audio_v3_tts")
        self.assertIn("not reachable", str(ctx.exception))
        self.assertIn("sgl-omni", str(ctx.exception))

    def test_booting_server_raises(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._connect([{"status": "unhealthy"}])
        self.assertIn("not healthy", str(ctx.exception))

    def test_foreign_hosted_model_raises_with_both_names(self):
        with self.assertRaises(RuntimeError) as ctx:
            self._connect([
                {"status": "healthy", "stages": []},
                {"data": [{"id": "Zyphra/zonos2"}]},
            ])
        message = str(ctx.exception)
        self.assertIn("Zyphra/zonos2", message)
        self.assertIn("bosonai/higgs-audio-v3-tts-4b", message)

    def test_matching_model_connects(self):
        client = self._connect([
            {"status": "healthy", "stages": []},
            {"data": [{"id": "bosonai/higgs-audio-v3-tts-4b"}]},
        ])
        self.assertEqual(client.entry.repo, "bosonai/higgs-audio-v3-tts-4b")


class PayloadTests(unittest.TestCase):
    """The /v1/audio/speech request shape per voice capability."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.ref = Path(self._tmp.name) / "narrator.wav"
        self.ref.write_bytes(b"abc")
        self.addCleanup(self._tmp.cleanup)

    def _make_client(self, model, **kwargs):
        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
        from backends.sglomni.catalog import entry_by_key
        client.entry = entry_by_key(model)
        client.api_url = "http://127.0.0.1:8100"
        client.voice = kwargs.get("voice")
        if "ref_audio" in kwargs:
            kwargs["ref_audio"] = str(self.ref)
        client.ref_audio = kwargs.get("ref_audio")
        client.ref_text = kwargs.get("ref_text", "")
        client.instructions = kwargs.get("instructions", "")
        client.language = "English"
        client._seed = None
        client._kv_fit = None
        return client

    def test_speaker_payload_sends_the_preset_name(self):
        client = self._make_client("qwen3_tts_0_6b_customvoice",
                                   voice="Vivian")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["voice"], "Vivian")
        self.assertEqual(payload["model"],
                         "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice")
        self.assertEqual(payload["response_format"], "wav")
        self.assertNotIn("ref_audio", payload)
        self.assertNotIn("task_type", payload)

    def test_speaker_without_voice_uses_the_server_default(self):
        client = self._make_client("voxtral_tts")
        self.assertEqual(client._request_payload("Hello.")["voice"],
                         "default")

    def test_design_payload_sends_task_type_and_instructions(self):
        client = self._make_client("qwen3_tts_1_7b_voicedesign",
                                   instructions="A warm narrator.")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["task_type"], "VoiceDesign")
        self.assertEqual(payload["instructions"], "A warm narrator.")

    def test_clone_payload_sends_reference_path_on_loopback(self):
        client = self._make_client("higgs_audio_v3_tts",
                                   ref_audio=str(self.ref),
                                   ref_text="A transcript.")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["ref_audio"], str(self.ref.resolve()))
        self.assertEqual(payload["ref_text"], "A transcript.")

    def test_clone_payload_inlines_audio_for_remote_servers(self):
        client = self._make_client("higgs_audio_v3_tts",
                                   ref_audio=str(self.ref))
        client.api_url = "http://10.20.30.40:8100"
        payload = client._request_payload("Hello.")
        self.assertTrue(payload["ref_audio"].startswith(
            "data:audio/wav;base64,"))
        self.assertEqual(
            base64.b64decode(payload["ref_audio"].partition(";base64,")[2]),
            b"abc")
        self.assertNotIn("ref_text", payload)

    def test_clone_without_reference_sends_no_reference_fields(self):
        client = self._make_client("higgs_audio_v3_tts")
        payload = client._request_payload("Hello.")
        self.assertNotIn("ref_audio", payload)
        self.assertEqual(payload["voice"], "default")

    def test_seed_included_when_resolved(self):
        client = self._make_client("qwen3_tts_1_7b_base",
                                   ref_audio="x.wav")
        client._seed = 7
        self.assertEqual(client._request_payload("Hello.")["seed"], 7)

    def test_zonos2_payload_raises_the_generation_cap(self):
        """Zonos2's 1024-frame engine default caps a request at ~12 s."""
        client = self._make_client("zonos2")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["max_new_tokens"], 12288)

    def test_higgs_payload_raises_the_generation_cap(self):
        """Higgs's 2048-frame engine default caps a request at ~27 s
        (75 fps); the catalog raises it to the most its admission window
        allows (~40 s after the prompt tokens)."""
        client = self._make_client("higgs_audio_v3_tts")
        payload = client._request_payload("Hello.")
        self.assertEqual(payload["max_new_tokens"], 3000)

    def test_payload_keeps_a_learned_kv_fit(self):
        """A capacity learned from an admission rejection caps later
        requests below the catalog value."""
        client = self._make_client("higgs_audio_v3_tts")
        client._kv_fit = 2500
        self.assertEqual(client._request_payload("Hello.")["max_new_tokens"],
                         2500)

    def test_models_without_a_cap_send_no_max_new_tokens(self):
        client = self._make_client("moss_tts")
        self.assertNotIn("max_new_tokens",
                         client._request_payload("Hello."))


class RequestErrorTests(unittest.TestCase):
    """OpenAI-style error envelopes decide retryability."""

    def _make_client(self):
        return SgOmniTTSClient.__new__(SgOmniTTSClient)

    def test_bad_request_envelope_is_not_retryable(self):
        client = self._make_client()
        detail = json.dumps({"error": {
            "message": "voice 'nope' not found",
            "type": "BadRequestError", "code": 400}})
        error = client._request_error(400, detail)
        self.assertIsInstance(error, NonRetryableTTSError)
        self.assertIn("voice 'nope' not found", str(error))

    def test_server_error_is_retryable(self):
        client = self._make_client()
        error = client._request_error(503, "overloaded")
        self.assertNotIsInstance(error, NonRetryableTTSError)

    def test_non_json_4xx_is_not_retryable(self):
        client = self._make_client()
        error = client._request_error(422, "plain text rejection")
        self.assertIsInstance(error, NonRetryableTTSError)


_KV_REJECTION_BODY = json.dumps({"error": {
    "message": "Request requires more tokens than the thinker KV cache "
               "can hold (input_tokens=684, max_new_tokens=12288, "
               "required_tokens=12972, kv_capacity=4095). Current "
               "mem_fraction_static is 0.800; try setting "
               "--thinker-mem-fraction-static higher.",
    "type": "InternalServerError", "code": 500}})


def _http_error(code: int, body: str) -> urllib.error.HTTPError:
    return urllib.error.HTTPError(
        "http://127.0.0.1:8100/v1/audio/speech", code, "error",
        hdrs=None, fp=io.BytesIO(body.encode("utf-8")))


def _speech_response():
    response = MagicMock()
    response.__enter__.return_value = response
    response.read.return_value = _WAV_BYTES
    return response


class KvAdmissionTests(unittest.TestCase):
    """The KV-window admission rejection refits max_new_tokens once."""

    def _client(self):
        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
        from backends.sglomni.catalog import entry_by_key
        client.entry = entry_by_key("higgs_audio_v3_tts")
        client.api_url = "http://127.0.0.1:8100"
        client.voice = None
        client.ref_audio = None
        client.ref_text = ""
        client.instructions = ""
        client.language = "English"
        client._seed = None
        client.chunk_size = None
        client._kv_fit = None
        return client

    def test_fit_is_parsed_from_the_server_message(self):
        fit = self._client()._kv_admission_fit(_KV_REJECTION_BODY)
        # kv_capacity 4095 - input 684 - the 64-frame margin.
        self.assertEqual(fit, 3347)

    def test_fit_is_cached_for_later_requests(self):
        client = self._client()
        client._kv_admission_fit(_KV_REJECTION_BODY)
        client._kv_admission_fit(_KV_REJECTION_BODY)
        self.assertEqual(client._kv_fit, 3347)

    def test_unrelated_errors_do_not_fit(self):
        client = self._client()
        self.assertIsNone(client._kv_admission_fit("CUDA out of memory"))
        self.assertIsNone(client._kv_fit)

    def test_a_window_below_the_floor_raises_with_guidance(self):
        body = json.dumps({"error": {"message":
            "Request requires more tokens than the thinker KV cache can "
            "hold (input_tokens=4000, max_new_tokens=12288, "
            "required_tokens=16288, kv_capacity=4095).", "code": 500}})
        with self.assertRaises(NonRetryableTTSError) as ctx:
            self._client()._kv_admission_fit(body)
        self.assertIn("shorter reference clip", str(ctx.exception))

    def test_request_wav_refits_and_resends_once(self):
        client = self._client()
        with patch(
                "converter.clients.sglomni.urllib.request.urlopen",
                side_effect=[_http_error(500, _KV_REJECTION_BODY),
                             _speech_response()]) as mock_open:
            wav = client._request_wav("Hello.")
        self.assertEqual(wav, _WAV_BYTES)
        self.assertEqual(mock_open.call_count, 2)
        refit = json.loads(mock_open.call_args[0][0].data)
        self.assertEqual(refit["max_new_tokens"], 3347)

    def test_request_wav_surfaces_a_refit_that_fails_again(self):
        client = self._client()
        with patch(
                "converter.clients.sglomni.urllib.request.urlopen",
                side_effect=[_http_error(500, _KV_REJECTION_BODY),
                             _http_error(500, _KV_REJECTION_BODY)]):
            with self.assertRaises(RuntimeError) as ctx:
                client._request_wav("Hello.")
        message = str(ctx.exception)
        self.assertIn("HTTP 500", message)
        self.assertIn("thinker KV cache", message)
        self.assertNotIsInstance(ctx.exception, NonRetryableTTSError)

    def test_request_wav_does_not_refit_other_errors(self):
        client = self._client()
        with patch(
                "converter.clients.sglomni.urllib.request.urlopen",
                side_effect=[_http_error(500, "CUDA out of memory")]):
            with self.assertRaises(RuntimeError) as ctx:
                client._request_wav("Hello.")
        self.assertIn("CUDA out of memory", str(ctx.exception))
        self.assertIsNone(client._kv_fit)

    def test_request_wav_keeps_the_refit_across_sub_requests(self):
        # A tight window (a long reference clip): the fit binds below the
        # 3000-frame catalog cap, and every later request carries it.
        tight_body = json.dumps({"error": {"message":
            "Request requires more tokens than the thinker KV cache can "
            "hold (input_tokens=1500, max_new_tokens=3000, "
            "required_tokens=4500, kv_capacity=4095).", "code": 500}})
        client = self._client()
        with patch(
                "converter.clients.sglomni.urllib.request.urlopen",
                side_effect=[_http_error(500, tight_body),
                             _speech_response(),
                             _speech_response()]) as mock_open:
            client._request_wav("Hello.")
            client._request_wav("Hello again.")
        self.assertEqual(mock_open.call_count, 3)
        second = json.loads(mock_open.call_args[0][0].data)
        self.assertEqual(second["max_new_tokens"], 2531)


class GenerateChunkTests(unittest.TestCase):
    """Chunk generation: WAV output, sub-chunking, bookkeeping."""

    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self._sleep = patch("converter.clients.base.time.sleep")
        self._sleep.start()
        self.addCleanup(self._sleep.stop)
        self.addCleanup(self._tmp.cleanup)

    def _make_client(self):
        client = SgOmniTTSClient.__new__(SgOmniTTSClient)
        from backends.sglomni.catalog import entry_by_key
        client.entry = entry_by_key("higgs_audio_v3_tts")
        client.chunks_dir = Path(self._tmp.name)
        client.api_url = "http://127.0.0.1:8100"
        client.voice = None
        client.ref_audio = None
        client.ref_text = ""
        client.instructions = ""
        client.language = "English"
        client._seed = None
        client.chunk_size = None
        client._kv_fit = None
        return client

    def _read_wav(self, path):
        with wave.open(str(path), "rb") as wav_file:
            return wav_file.readframes(wav_file.getnframes())

    def test_generate_chunk_writes_the_wav_response(self):
        client = self._make_client()
        with patch.object(client, "_request_wav", return_value=_WAV_BYTES):
            result = client.generate_chunk("Hello world.", 1)
        self.assertIsNotNone(result)
        path = Path(result)
        self.assertEqual(path.name, "chunk_0001.wav")
        self.assertEqual(self._read_wav(path), _WAV_FRAMES)

    def test_long_text_is_subchunked_and_concatenated(self):
        client = self._make_client()
        text = " ".join(f"word{i}" for i in range(24))
        responses = [_WAV_BYTES, _WAV_BYTES, _WAV_BYTES]
        with patch.object(config, "CHUNK_SIZE", 10), \
                patch.object(client, "_request_wav",
                             side_effect=responses) as mock_wav, \
                patch("converter.clients.sglomni.concat_audio_files") as mock_concat:
            result = client.generate_chunk(text, 1)
        # 24 words at CHUNK_SIZE 10 -> three sub-requests (10/10/4).
        self.assertEqual(mock_wav.call_count, 3)
        self.assertIsNotNone(result)
        mock_concat.assert_called_once()
        args = mock_concat.call_args[0]
        self.assertEqual(len(args[0]), 3)
        self.assertEqual(args[1], Path(result))

    def test_run_chunk_size_caps_the_sub_requests(self):
        """The pre-flight clamp (a chunk_words-capped model's popup
        answer) overrides CHUNK_SIZE for this run."""
        client = self._make_client()
        client.chunk_size = 10
        text = " ".join(f"word{i}" for i in range(24))
        with patch.object(client, "_request_wav",
                          return_value=_WAV_BYTES) as mock_wav, \
                patch("converter.clients.sglomni.concat_audio_files"):
            client.generate_chunk(text, 1)
        self.assertEqual(mock_wav.call_count, 3)

    def test_single_subchunk_skips_concatenation(self):
        client = self._make_client()
        with patch.object(client, "_request_wav", return_value=_WAV_BYTES), \
                patch("converter.clients.sglomni.concat_audio_files") as mock_concat:
            client.generate_chunk("Hello.", 1)
        mock_concat.assert_not_called()

    def test_empty_text_fails_the_chunk(self):
        client = self._make_client()
        with patch.object(client, "_request_wav") as mock_wav:
            self.assertIsNone(client.generate_chunk("   ", 1))
        mock_wav.assert_not_called()

    def test_request_failure_fails_the_chunk_attempt(self):
        client = self._make_client()
        with patch.object(client, "_request_wav",
                          side_effect=RuntimeError("down")) as mock_wav:
            self.assertIsNone(client.generate_chunk("Hello.", 1))
        self.assertEqual(mock_wav.call_count, 1)

    def test_stale_chunk_files_are_removed(self):
        stale = Path(self._tmp.name) / "chunk_0001.mp3"
        stale.write_bytes(b"old")
        client = self._make_client()
        with patch.object(client, "_request_wav", return_value=_WAV_BYTES):
            client.generate_chunk("Hello.", 1)
        remaining = sorted(path.name for path in
                           Path(self._tmp.name).glob("chunk_0001.*"))
        self.assertEqual(remaining, ["chunk_0001.wav"])


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