VoiceStudio/tests/test_voxcpm2_subprocess.py

"""VoxCPM2 from its own venv: the sidecar, and the switch to it.

The sidecar runs in a venv that holds only voxcpm and its dependencies, so it
must import nothing from the app, and it must drive the model exactly as the
in-process VoxCPM2Backend does. These tests run it with a fake `voxcpm`.
"""
import importlib.util
import io
import json
import re
import struct
import sys
import types
from pathlib import Path

import numpy as np
import pytest

_MAIN = Path(__file__).resolve().parents[1] / "backend/engines/voxcpm2_subprocess/main.py"


def _load_sidecar(monkeypatch, calls):
    class FakeModel:
        sample_rate = 48000

        def generate(self, **kw):
            calls.append(kw)
            return np.zeros(480, dtype=np.float32)

    class VoxCPM:
        @classmethod
        def from_pretrained(cls, checkpoint, **kw):
            calls.append({"from_pretrained": checkpoint, **kw})
            return FakeModel()

    fake = types.ModuleType("voxcpm")
    fake.VoxCPM = VoxCPM
    monkeypatch.setitem(sys.modules, "voxcpm", fake)
    spec = importlib.util.spec_from_file_location("_voxcpm2_sidecar_under_test", _MAIN)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def _frames(buf):
    data, out, i = buf.getvalue(), [], 0
    while i < len(data):
        (n,) = struct.unpack("!I", data[i:i + 4])
        out.append(json.loads(data[i + 4:i + 4 + n]))
        i += 4 + n
    return out


def test_voice_design_maps_to_voice_description(monkeypatch):
    calls = []
    sidecar = _load_sidecar(monkeypatch, calls)
    out = io.BytesIO()

    sidecar._handle_synthesize({"op": "synthesize", "text": "hello", "description": "warm, low"}, out)

    assert calls[-1] == {
        "text": "hello",
        "voice_description": "warm, low",
        "cfg_value": 2.0,
        "inference_timesteps": 10,
    }
    audio = _frames(out)[-1]
    assert audio["op"] == "audio"
    assert audio["sample_rate"] == 48000 and audio["n_samples"] == 480


def test_clone_mode_passes_the_reference_and_prompt(monkeypatch, tmp_path):
    calls = []
    sidecar = _load_sidecar(monkeypatch, calls)
    ref = tmp_path / "ref.wav"
    ref.write_bytes(b"x")

    sidecar._handle_synthesize(
        {"text": "hi", "ref_audio": str(ref), "ref_text": "hello there", "instruct": "calm",
         "guidance_scale": 3.0, "num_step": 12},
        io.BytesIO(),
    )

    assert calls[-1] == {
        "text": "(calm)hi",
        "cfg_value": 3.0,
        "inference_timesteps": 12,
        "reference_wav_path": str(ref),
        "prompt_wav_path": str(ref),
        "prompt_text": "hello there",
    }


def test_a_reference_without_its_transcript_is_not_a_prompt(monkeypatch, tmp_path):
    calls = []
    sidecar = _load_sidecar(monkeypatch, calls)
    ref = tmp_path / "ref.wav"
    ref.write_bytes(b"x")
    sidecar._handle_synthesize({"text": "hi", "ref_audio": str(ref)}, io.BytesIO())
    assert calls[-1]["reference_wav_path"] == str(ref)
    assert calls[-1]["prompt_wav_path"] is None


def test_loads_the_configured_checkpoint_without_the_denoiser(monkeypatch):
    calls = []
    monkeypatch.setenv("OMNIVOICE_VOXCPM_MODEL", "local/ckpt")
    sidecar = _load_sidecar(monkeypatch, calls)
    sidecar._handle_synthesize({"text": "hi"}, io.BytesIO())
    assert calls[0] == {"from_pretrained": "local/ckpt", "load_denoiser": False}


def test_rejects_a_url_reference(monkeypatch):
    sidecar = _load_sidecar(monkeypatch, [])
    with pytest.raises(ValueError, match="local file path"):
        sidecar._handle_synthesize({"text": "hi", "ref_audio": "https://x.test/y.wav"}, io.BytesIO())


def test_retries_a_transient_download_failure_only(monkeypatch):
    sidecar = _load_sidecar(monkeypatch, [])
    monkeypatch.setattr(sidecar.time, "sleep", lambda s: None)
    attempts = []

    def flaky():
        attempts.append(1)
        if len(attempts) < 3:
            raise OSError("peer closed connection without sending complete message body")
        return "model"

    assert sidecar._with_retries(flaky) == "model"
    assert len(attempts) == 3

    broken_calls = []

    def broken():
        broken_calls.append(1)
        raise ValueError("bad config")

    with pytest.raises(ValueError):
        sidecar._with_retries(broken)
    assert len(broken_calls) == 1  # a permanent error is not retried


def test_the_sidecar_imports_nothing_from_the_app():
    src = _MAIN.read_text(encoding="utf-8")
    for name in ("services", "core", "engines", "api", "backend", "utils"):
        assert not re.search(rf"^\s*(from|import) {name}\b", src, re.M), name


def test_the_class_switches_to_the_sidecar_once_its_venv_exists(monkeypatch, tmp_path):
    from engines.voxcpm2_subprocess import VoxCPM2SubprocessBackend
    from services import tts_backend
    from services.sidecar_install import _INSTALL_COMPLETE_MARKER, _venv_python

    monkeypatch.setenv("OMNIVOICE_VOXCPM2_DIR", "")
    monkeypatch.delenv("OMNIVOICE_VOXCPM2_DIR")
    assert tts_backend.get_backend_class("voxcpm2") is tts_backend.VoxCPM2Backend

    py = _venv_python(tmp_path / ".venv")
    py.parent.mkdir(parents=True)
    py.write_text("#!fake\n")
    (tmp_path / _INSTALL_COMPLETE_MARKER).write_text("x\n", encoding="utf-8")
    monkeypatch.setenv("OMNIVOICE_VOXCPM2_DIR", str(tmp_path))

    cls = tts_backend.get_backend_class("voxcpm2")
    assert cls is VoxCPM2SubprocessBackend
    assert cls.venv_python() == py
    assert cls.is_available() == (True, "ready")
    # The same engine to the rest of the app.
    for attr in ("id", "display_name", "supports_voice_design", "applies_own_mastering", "gpu_compat"):
        assert getattr(cls, attr) == getattr(tts_backend.VoxCPM2Backend, attr), attr
    assert cls().supported_languages == tts_backend.VoxCPM2Backend().supported_languages


def test_the_sidecar_class_prepares_the_reference_and_trims_the_tail(monkeypatch):
    import torch

    import services.audio_dsp as dsp
    from engines.voxcpm2_subprocess import VoxCPM2SubprocessBackend
    from services import tts_backend

    sent = {}

    def fake_generate(self, text, **kw):
        sent.update(kw)
        return torch.zeros(1, 4800)

    trimmed = {}

    def fake_trim(wav, sr):
        trimmed["sr"] = sr
        return wav[:, :10]

    # Patch the class the engine actually inherits from: other suites purge
    # `services` from sys.modules, so a fresh import of subprocess_backend
    # can be a different module than the one the engine subclassed.
    monkeypatch.setattr(VoxCPM2SubprocessBackend.__bases__[0], "generate", fake_generate)
    monkeypatch.setattr(tts_backend, "_prepare_voxcpm_ref", lambda p: p + ".prepared.wav")
    monkeypatch.setattr(dsp, "trim_trailing_silence", fake_trim)

    out = VoxCPM2SubprocessBackend().generate("hi", ref_audio="/clip.wav")

    assert sent["ref_audio"] == "/clip.wav.prepared.wav"
    assert trimmed["sr"] == 48000
    assert tuple(out.shape) == (1, 10)


def test_output_is_resampled_to_the_48_khz_the_parent_assumes(monkeypatch):
    """The parent trims and labels the PCM at a fixed 48 kHz, so a model that
    reports another rate must be resampled, not passed through."""
    sidecar = _load_sidecar(monkeypatch, [])
    sidecar._handle_synthesize({"text": "warm up"}, io.BytesIO())
    type(sidecar._MODEL).sample_rate = 24000
    out = io.BytesIO()
    sidecar._handle_synthesize({"text": "hi"}, out)
    audio = _frames(out)[-1]
    assert audio["sample_rate"] == 48000
    assert audio["n_samples"] == 960  # 480 samples at 24 kHz


def test_a_failed_install_leaves_voxcpm2_in_process(monkeypatch, tmp_path):
    """A venv interpreter without the completion marker (a reinstall that
    failed partway) must not hide the working in-process engine."""
    from services import tts_backend
    from services.sidecar_install import _venv_python

    py = _venv_python(tmp_path / ".venv")
    py.parent.mkdir(parents=True)
    py.write_text("#!fake\n")
    monkeypatch.setenv("OMNIVOICE_VOXCPM2_DIR", str(tmp_path))
    assert tts_backend.get_backend_class("voxcpm2") is tts_backend.VoxCPM2Backend