"""Tests for Fish Speech DAC-code caching via VoiceEmbeddingCache. Covers: - Cache miss → DAC encode → store - Cache hit → skip DAC encode, reuse cached ref_codes_fq - Inline ref_audio (no voice name) → no caching, full encode path - Stale-cache protection via created_at - Temp file cleanup on cache hit """ import os import tempfile from pathlib import Path import numpy as np import pytest import torch from pytest_mock import MockerFixture pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _make_info_dict( *, text: str = "Hello world", ref_text: str = "Reference transcript", ref_audio_sr: int = 44100, voice_name: str | None = None, voice_created_at: float | None = None, ref_audio_path: str | None = None, ) -> dict: """Build a minimal info_dict for _build_structured_voice_clone_prefill_embeds.""" d: dict = { "text": text, "ref_text": ref_text, "ref_audio_sr": ref_audio_sr, "fish_structured_voice_clone": True, } if ref_audio_path is not None: d["ref_audio_path"] = ref_audio_path if voice_name is not None: d["voice_name"] = voice_name if voice_created_at is not None: d["voice_created_at"] = voice_created_at return d def _write_temp_npy(wav: np.ndarray | None = None) -> str: """Write a temporary .npy file with dummy audio and return its path.""" if wav is None: wav = np.random.randn(44100).astype(np.float32) # 1 second @ 44.1kHz with tempfile.NamedTemporaryFile(prefix="fish_test_", suffix=".npy", delete=False) as f: np.save(f, wav) return f.name # Fake ref_codes_fq: [frames, codebooks] _FAKE_REF_CODES = torch.randint(0, 1024, (10, 10), dtype=torch.long) class TestFishSpeechVoiceCacheIntegration: """Test the cache-hit / cache-miss / no-cache paths in the model.""" @pytest.fixture def mock_model(self, mocker: MockerFixture): """Create a mock FishSpeechSlowARForConditionalGeneration with cache.""" from vllm_omni.utils.voice_cache import VoiceEmbeddingCache model = mocker.MagicMock() model._voice_cache = VoiceEmbeddingCache(max_entries=4) model._semantic_begin_id = 151678 model._num_codebooks = 10 model._codebook_size = 4096 model.model_path = "/fake/model" model.codebook_embeddings = mocker.MagicMock() model.codebook_embeddings.weight = mocker.MagicMock() model.codebook_embeddings.weight.device = torch.device("cpu") return model def test_cache_miss_stores_codes(self, mock_model): """First request with a named voice should encode and store in cache.""" cache = mock_model._voice_cache voice_name = "alice" created_at = 1712345678.0 # Verify cache starts empty. key = cache.make_cache_key(voice_name, xvec_only=False, created_at=created_at) assert cache.get(key) is None # Simulate a cache store (what the model does on miss). cache.put(key, {"ref_codes_fq": _FAKE_REF_CODES.detach().cpu()}) # Verify it's now cached. cached = cache.get(key) assert cached is not None assert torch.equal(cached["ref_codes_fq"], _FAKE_REF_CODES) def test_cache_hit_returns_cached_codes(self, mock_model): """Second request with same voice should hit cache.""" cache = mock_model._voice_cache voice_name = "alice" created_at = 1712345678.0 key = cache.make_cache_key(voice_name, xvec_only=False, created_at=created_at) cache.put(key, {"ref_codes_fq": _FAKE_REF_CODES.detach().cpu()}) # Hit. cached = cache.get(key) assert cached is not None ref_codes = cached["ref_codes_fq"].to(device=torch.device("cpu"), dtype=torch.long) assert torch.equal(ref_codes, _FAKE_REF_CODES) assert cache.stats()["hits"] >= 1 def test_no_voice_name_skips_cache(self, mock_model): """Inline ref_audio without voice_name should not use cache.""" cache = mock_model._voice_cache # Without voice_name, the model should not interact with cache at all. info = _make_info_dict(voice_name=None, ref_audio_path=_write_temp_npy()) assert info.get("voice_name") is None # Cache should remain untouched. assert cache.stats()["hits"] == 0 assert cache.stats()["misses"] == 0 def test_stale_cache_on_reupload(self, mock_model): """Re-uploading a voice (new created_at) should not hit old cache.""" cache = mock_model._voice_cache voice_name = "alice" key_old = cache.make_cache_key(voice_name, xvec_only=False, created_at=1000.0) cache.put(key_old, {"ref_codes_fq": _FAKE_REF_CODES}) # Re-upload produces a different created_at. key_new = cache.make_cache_key(voice_name, xvec_only=False, created_at=2000.0) assert cache.get(key_new) is None # miss assert cache.get(key_old) is not None # old still there def test_temp_file_cleaned_on_cache_hit(self): """On cache hit, the temp .npy file written by the entrypoint should be deleted.""" tmp_path = _write_temp_npy() assert os.path.exists(tmp_path) # Simulate what the model does on cache hit: remove the temp file. try: os.remove(tmp_path) except OSError: pass assert not os.path.exists(tmp_path) def test_created_at_zero_disables_cache(self, mock_model): """created_at=0 should not create a cache key (caching disabled).""" cache = mock_model._voice_cache info = _make_info_dict( voice_name="bob", voice_created_at=0.0, ref_audio_path=_write_temp_npy(), ) # The model checks: if _created_at > 0 → enable cache. # With 0.0, no cache interaction should happen. _created_at = float(info.get("voice_created_at", 0)) assert _created_at <= 0 assert cache.stats()["hits"] == 0 assert cache.stats()["misses"] == 0 class TestFishSpeechValidatorUploadedVoice: """Test _validate_fish_tts_request uploaded voice resolution.""" def test_uploaded_voice_resolves_ref_audio( self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, ): """When voice matches an uploaded speaker, ref_audio should be auto-set.""" request = mocker.MagicMock() request.input = "Hello" request.voice = "alice" request.ref_audio = None request.ref_text = None request.max_new_tokens = None # Uploaded speaker with ref_text. uploaded_speakers = { "alice": { "file_path": "/tmp/fake_audio.wav", "ref_text": "Hi this is Alice", "created_at": 1712345678, }, } # Simulate: voice in uploaded_speakers, file exists, get_audio returns data URL. monkeypatch.setattr(Path, "exists", lambda self: True) voice_lower = request.voice.lower() assert voice_lower in uploaded_speakers speaker_info = uploaded_speakers[voice_lower] ref_text_from_upload = speaker_info.get("ref_text") assert ref_text_from_upload == "Hi this is Alice" def test_uploaded_voice_without_ref_text_uses_request_ref_text( self, mocker: MockerFixture, ): """If upload has no ref_text but request provides it, use request's.""" request = mocker.MagicMock() request.input = "Hello" request.voice = "bob" request.ref_audio = None request.ref_text = "Request-level transcript" request.max_new_tokens = None uploaded_speakers = { "bob": { "file_path": "/tmp/fake_audio.wav", "ref_text": None, "created_at": 1712345678, }, } voice_lower = request.voice.lower() speaker_info = uploaded_speakers[voice_lower] upload_ref_text = speaker_info.get("ref_text") # Upload has no ref_text, so request.ref_text should remain. assert upload_ref_text is None assert request.ref_text == "Request-level transcript"