tests/test_qwentts_cpp_backend.py

"""Tests for the qwentts.cpp in-process backend.

The C library is mocked via a fake QwenttsModel; we don't load any
real models. Voice library reads use the per-test isolated voices
dir set by the autouse fixture.
"""

from __future__ import annotations

import asyncio
import io
import os
import wave
from pathlib import Path
from typing import Any
from unittest import mock

import numpy as np
import pytest


from reachy_conv_app.tts import qwentts_cpp_backend  # noqa: E402
from reachy_conv_app.tts import voice_library        # noqa: E402
from reachy_conv_app.tts.base import TTSBackend, VoiceSpec  # noqa: E402


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture(autouse=True)
def _stub_cpu_preflight(monkeypatch):
    """The CPU pre-flight checks /proc/cpuinfo for AMX. This dev box
    doesn't have AMX (it would segfault on the real C library too),
    so we patch the check to return supported. Tests that *want* the
    pre-flight to fail can re-patch inside the test body."""
    from reachy_conv_app.tts import qwentts_variant
    monkeypatch.setattr(
        qwentts_variant, "cpu_supports_qwentts",
        lambda: (True, ""),
    )

class _FakeSynthResult:
    def __init__(self, samples: np.ndarray, sample_rate: int):
        self.samples = samples
        self.sample_rate = sample_rate


class _FakeModel:
    """A drop-in replacement for QwenttsModel."""

    def __init__(self):
        self.is_loaded = True
        self._loaded_voice_mode = "preset"
        self._loaded_model_size = "0.6B"
        self._status = mock.MagicMock()
        self._status.state = "ready"
        # `status` is the property the real model exposes
        self.status = self._status
        self._speakers = ["Vivian", "Serena", "Uncle_Fu"]
        self._synth_log: list[dict] = []
        self._speakers_should_raise = False
        self._force_unloaded = False

    def get_speakers(self):
        if self._speakers_should_raise:
            return []
        return list(self._speakers)

    def load(self, voice_mode, model_size, *, progress=None):
        self._loaded_voice_mode = voice_mode
        self._loaded_model_size = model_size
        self.is_loaded = True
        self._status.state = "ready"
        return self._status

    def close(self):
        self.is_loaded = False
        self._status.state = "unloaded"

    def synth(self, req):
        self._synth_log.append({"text": req.text, "voice_mode": req.voice_mode,
                                "speaker": req.speaker, "voice_name": req.voice_name,
                                "voice_description": req.voice_description,
                                "language": req.language, "instruct": req.instruct,
                                "seed": req.seed, "temperature": req.temperature})
        n = 2400  # 0.1s at 24 kHz
        return _FakeSynthResult(np.zeros(n, dtype=np.float32), 24000)

    async def synth_async(self, req):
        return self.synth(req)


@pytest.fixture(autouse=True)
def isolated_voices_dir(tmp_path, monkeypatch):
    # v0.5.25: also isolate chatterbox_voices_root so a stray
    # chatterbox voice dir under the real ~/.cache/... doesn't
    # pollute the qwentts list_voices assertions.
    monkeypatch.setenv("REACHY_CONV_APP_VOICES_DIR", str(tmp_path / "voices"))
    monkeypatch.setenv(
        "REACHY_CONV_APP_CHATTERBOX_VOICES_DIR", str(tmp_path / "cb_voices")
    )
    yield


def _make_backend(model: _FakeModel | None = None) -> qwentts_cpp_backend.QwenTtsCpp:
    return qwentts_cpp_backend.QwenTtsCpp(model=model or _FakeModel())


# ---------------------------------------------------------------------------
# Basic plumbing
# ---------------------------------------------------------------------------

def test_backend_name():
    b = _make_backend()
    assert b.name == "qwentts_cpp"


def test_backend_is_subclass_of_tts_backend():
    assert issubclass(qwentts_cpp_backend.QwenTtsCpp, TTSBackend)


def test_is_available_when_loaded():
    model = _FakeModel()
    model.is_loaded = True
    model._status.state = "ready"
    b = qwentts_cpp_backend.QwenTtsCpp(model=model)
    assert b.is_available() is True


def test_is_available_false_when_wheel_missing(monkeypatch):
    """is_available() should return False if the qwentts_cpp wheel can't import."""
    import builtins
    real_import = builtins.__import__

    def fake_import(name, *args, **kwargs):
        if name == "qwentts_cpp" or name.startswith("qwentts_cpp."):
            raise ImportError("simulated wheel not installed")
        return real_import(name, *args, **kwargs)

    monkeypatch.setattr(builtins, "__import__", fake_import)
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model)
    assert b.is_available() is False
    # An availability_error should be set so the dashboard can show why
    assert b.availability_error is not None
    assert "qwentts_cpp" in b.availability_error


def test_is_available_true_when_wheel_present(monkeypatch):
    """is_available() should return True if qwentts_cpp imports cleanly."""
    # The real wheel IS importable in this venv, so just call it.
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model)
    # Don't call _ensure_loaded — the constructor doesn't load any more.
    assert b.is_available() is True


# ---------------------------------------------------------------------------
# list_voices
# ---------------------------------------------------------------------------

def test_list_voices_preset_returns_speakers():
    model = _FakeModel()
    model._speakers = ["Vivian", "Serena"]
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset")
    voices = b.list_voices()
    names = [v.id for v in voices]
    assert names == ["Vivian", "Serena"]
    assert all(isinstance(v, VoiceSpec) for v in voices)


def test_list_voices_cloned_returns_library_entries():
    spk = np.zeros(64, dtype=np.float32)
    codes = np.zeros((4, 10), dtype=np.int32)
    voice_library.add_voice(
        name="alice", ref_text="x",
        spk_embedding=spk, rvq_codes=codes,
        source_wav=b"RIFF" + b"\x00" * 100,
        model_id="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
    )
    voice_library.add_voice(
        name="bob", ref_text="y",
        spk_embedding=spk, rvq_codes=codes,
        source_wav=b"RIFF" + b"\x00" * 100,
        model_id="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
    )
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="cloned")
    voices = b.list_voices()
    assert sorted(v.id for v in voices) == ["alice", "bob"]


def test_list_voices_design_returns_empty():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="design")
    assert b.list_voices() == []


# ---------------------------------------------------------------------------
# synth — preset
# ---------------------------------------------------------------------------

def test_synth_preset_uses_speaker_from_voice_arg():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset")
    pcm, sr = asyncio.run(b.synth("Hello.", "Serena"))
    assert sr == 24000
    # PCM int16 array, non-empty
    assert isinstance(pcm, np.ndarray)
    assert pcm.dtype == np.int16
    assert pcm.size > 0
    # The synth call recorded the right speaker
    assert model._synth_log[-1]["speaker"] == "Serena"
    assert model._synth_log[-1]["voice_mode"] == "preset"


def test_synth_preset_defaults_to_vivian_when_voice_empty():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset")
    asyncio.run(b.synth("Hi", ""))
    assert model._synth_log[-1]["speaker"] == "Vivian"


# ---------------------------------------------------------------------------
# synth — cloned
# ---------------------------------------------------------------------------

def test_synth_cloned_uses_voice_name():
    spk = np.zeros(64, dtype=np.float32)
    codes = np.zeros((16, 50), dtype=np.int32)
    voice_library.add_voice(
        name="alice", ref_text="x",
        spk_embedding=spk, rvq_codes=codes,
        source_wav=b"RIFF" + b"\x00" * 100,
        model_id="Qwen/Qwen3-TTS-12Hz-0.6B-Base",
    )
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="cloned")
    asyncio.run(b.synth("Hi", "alice"))
    assert model._synth_log[-1]["voice_name"] == "alice"
    assert model._synth_log[-1]["voice_mode"] == "cloned"


# ---------------------------------------------------------------------------
# synth — design
# ---------------------------------------------------------------------------

def test_synth_design_uses_voice_as_description():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="design")
    asyncio.run(b.synth("Hi", "a deep British voice"))
    assert model._synth_log[-1]["voice_description"] == "a deep British voice"


# ---------------------------------------------------------------------------
# advanced controls
# ---------------------------------------------------------------------------

def test_synth_passes_instruct():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(
        model=model, voice_mode="preset", instruct="speak slowly and warmly",
    )
    asyncio.run(b.synth("Hi", "Vivian"))
    assert model._synth_log[-1]["instruct"] == "speak slowly and warmly"


def test_synth_passes_seed_and_temperature():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(
        model=model, voice_mode="preset", seed=42, temperature=0.3,
    )
    asyncio.run(b.synth("Hi", "Vivian"))
    assert model._synth_log[-1]["seed"] == 42
    assert model._synth_log[-1]["temperature"] == 0.3


def test_synth_passes_language():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(
        model=model, voice_mode="preset", language="japanese",
    )
    asyncio.run(b.synth("Hi", "Vivian"))
    assert model._synth_log[-1]["language"] == "japanese"


def test_set_advanced_updates_knobs():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset", seed=1)
    b.set_advanced(seed=99, temperature=0.0, top_p=0.5)
    asyncio.run(b.synth("Hi", "Vivian"))
    log = model._synth_log[-1]
    assert log["seed"] == 99
    assert log["temperature"] == 0.0


# ---------------------------------------------------------------------------
# empty text
# ---------------------------------------------------------------------------

def test_synth_empty_text_returns_short_silence_pcm():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset")
    pcm, sr = asyncio.run(b.synth("", "Vivian"))
    assert sr == 24000
    # Should be a short silence buffer (no model call)
    assert isinstance(pcm, np.ndarray)
    assert pcm.dtype == np.int16
    assert pcm.size <= 240  # <= 10ms at 24kHz
    # No synth call should have been made
    assert model._synth_log == []


def test_synth_whitespace_only_text_returns_silence():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset")
    wav, sr = asyncio.run(b.synth("   \n  ", "Vivian"))
    assert sr == 24000
    # No synth call should have been made
    assert model._synth_log == []


# ---------------------------------------------------------------------------
# voice mode / size switching
# ---------------------------------------------------------------------------

def test_set_voice_mode_reloads_model():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset", model_size="0.6B")
    assert model._loaded_voice_mode == "preset"
    b.set_voice_mode("cloned", "0.6B")
    assert model._loaded_voice_mode == "cloned"
    b.set_voice_mode("design", "1.7B")
    assert model._loaded_voice_mode == "design"
    assert model._loaded_model_size == "1.7B"


def test_set_voice_mode_same_mode_no_reload(monkeypatch):
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model, voice_mode="preset", model_size="0.6B")
    load_calls = []
    orig = model.load
    def tracker(*a, **kw):
        load_calls.append((a, kw))
        return orig(*a, **kw)
    model.load = tracker
    b.set_voice_mode("preset", "0.6B")  # same
    assert load_calls == []


# ---------------------------------------------------------------------------
# close
# ---------------------------------------------------------------------------

def test_close_frees_model():
    model = _FakeModel()
    b = qwentts_cpp_backend.QwenTtsCpp(model=model)
    b.close()
    assert model.is_loaded is False