test_modules/test_tts_pregen.py

"""Dependency-light tests for TTS pregeneration behavior."""
import os
import logging
import sys
import tempfile
import time
import types
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
logging.disable(logging.CRITICAL)


class FakeArray(list):
    def __mul__(self, value):
        if isinstance(value, list):
            return FakeArray(a * b for a, b in zip(self, value))
        return FakeArray(item * value for item in self)

    __rmul__ = __mul__


def install_dependency_stubs():
    try:
        import numpy as real_np  # noqa: F401
    except ModuleNotFoundError:
        fake_np = types.ModuleType("numpy")
        fake_np.ndarray = FakeArray
        fake_np.float32 = float
        fake_np.pi = 3.141592653589793
        fake_np.ones = lambda n, dtype=None: FakeArray([1.0] * n)
        fake_np.zeros = lambda n, dtype=None: FakeArray([0.0] * n)
        fake_np.linspace = lambda start, stop, num, dtype=None: FakeArray(
            [start + (stop - start) * i / max(num - 1, 1) for i in range(num)]
        )
        fake_np.sin = lambda values: FakeArray(
            __import__("math").sin(value) for value in values
        )
        fake_np.concatenate = lambda arrays: FakeArray(
            item for array in arrays for item in array
        )
        sys.modules["numpy"] = fake_np

    try:
        import soundfile as real_sf  # noqa: F401
    except ModuleNotFoundError:
        fake_sf = types.ModuleType("soundfile")
        audio_store = {}

        def write(path, waveform, sample_rate):
            audio_store[str(path)] = (waveform, sample_rate)
            Path(path).write_bytes(b"RIFF fake wav")

        def read(path, dtype="float32"):
            waveform, sample_rate = audio_store[str(path)]
            return waveform, sample_rate

        fake_sf.write = write
        fake_sf.read = read
        sys.modules["soundfile"] = fake_sf


install_dependency_stubs()

import numpy as np  # noqa: E402
import tts  # noqa: E402


PASS = 0
FAIL = 0


def check(label, condition, detail=""):
    global PASS, FAIL
    if condition:
        PASS += 1
        print(f"[OK] {label}")
    else:
        FAIL += 1
        print(f"[FAIL] {label} -- {detail}")


def wait_for_pregen(chunks, profile_id=None, timeout=2.0):
    deadline = time.time() + timeout
    status = tts.get_pregeneration_status(chunks, profile_id)
    while status.get("in_progress") and time.time() < deadline:
        time.sleep(0.02)
        status = tts.get_pregeneration_status(chunks, profile_id)
    return status


def reset_pregen_state():
    tts.cancel_pregeneration()
    with tts._pregen_lock:
        tts._pregen_cache.clear()
        tts._pregen_in_progress.clear()
        tts._pregen_progress.clear()
        tts._pregen_cancel_events.clear()


def run_with_temp_cache(test_fn):
    original_cache_dir = tts._CACHE_DIR
    original_synth = tts._synthesize_single
    with tempfile.TemporaryDirectory() as tmp:
        tts._CACHE_DIR = tmp
        os.makedirs(tts._CACHE_DIR, exist_ok=True)
        reset_pregen_state()
        try:
            test_fn()
        finally:
            reset_pregen_state()
            tts._CACHE_DIR = original_cache_dir
            tts._synthesize_single = original_synth


def test_pregen_limits_initial_chunks():
    calls = []

    def fake_synth(chunk, profile_id):
        calls.append(chunk)
        return np.ones(8, dtype=np.float32), 24000

    tts._synthesize_single = fake_synth
    chunks = ["one", "two", "three", "four", "five"]
    tts.pregenerate_story_audio(chunks, voice_profile_id="voice-a", max_chunks=2)
    status = wait_for_pregen(chunks, "voice-a")

    check("pregen targets requested initial chunks", status.get("target") == 2, status)
    check("pregen synthesized only initial chunks", calls == chunks[:2], calls)
    check("partial pregen is not marked complete", status.get("complete") is False, status)
    check("partial pregen records cached count", status.get("cached") == 2, status)


def test_pregen_failure_records_error():
    def fake_synth(chunk, profile_id):
        if chunk == "bad":
            raise RuntimeError("boom")
        return np.ones(8, dtype=np.float32), 24000

    tts._synthesize_single = fake_synth
    chunks = ["ok", "bad", "later"]
    tts.pregenerate_story_audio(chunks, voice_profile_id="voice-b", max_chunks=3)
    status = wait_for_pregen(chunks, "voice-b")

    check("failed pregen records error", status.get("errors") == 1, status)
    check("failed pregen is not complete", status.get("complete") is False, status)
    check("failed chunk is not cached", tts._get_cached_audio("bad", "voice-b") is None)


def test_cancel_stops_waiting_pregen():
    calls = []

    def fake_synth(chunk, profile_id):
        calls.append(chunk)
        return np.ones(8, dtype=np.float32), 24000

    tts._synthesize_single = fake_synth
    chunks = ["slow-one", "slow-two"]
    tts.GPU_INFERENCE_LOCK.acquire()
    try:
        tts.pregenerate_story_audio(chunks, voice_profile_id="voice-c", max_chunks=2)
        time.sleep(0.05)
        tts.cancel_pregeneration()
    finally:
        tts.GPU_INFERENCE_LOCK.release()

    status = wait_for_pregen(chunks, "voice-c")
    time.sleep(0.05)
    check("cancel marks pregen cancelled", status.get("cancelled") is True, status)
    check("cancelled pregen does not synthesize queued chunks", calls == [], calls)


def test_progress_metadata_is_bounded():
    original_limit = tts._PREGEN_PROGRESS_MAX_STORIES
    tts._PREGEN_PROGRESS_MAX_STORIES = 2
    try:
        with tts._pregen_lock:
            for idx in range(5):
                tts._pregen_progress[f"story-{idx}"] = {
                    "cached": 0,
                    "total": 1,
                    "target": 1,
                    "in_progress": False,
                    "complete": False,
                    "errors": 0,
                    "cancelled": False,
                }
            tts._prune_pregen_progress_locked()
            keys = list(tts._pregen_progress)
        check("pregen progress metadata is bounded", len(keys) == 2, keys)
        check("pregen progress keeps newest records", keys == ["story-3", "story-4"], keys)
    finally:
        tts._PREGEN_PROGRESS_MAX_STORIES = original_limit


if __name__ == "__main__":
    run_with_temp_cache(test_pregen_limits_initial_chunks)
    run_with_temp_cache(test_pregen_failure_records_error)
    run_with_temp_cache(test_cancel_stops_waiting_pregen)
    run_with_temp_cache(test_progress_metadata_is_bounded)
    print(f"TTS pregen tests: {PASS} passed, {FAIL} failed")
    sys.exit(1 if FAIL else 0)