UmbraAudio/tests/test_variations.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

import io

import httpx
import numpy as np
import soundfile as sf

from umbraaudio.generation.chapter import (
    EMOTION_INSTRUCT_MAP,
    EMOTION_OPTIONS,
    NARRATOR_COMPRESSION_INSTRUCT,
    PACE_INSTRUCT,
    PACE_OPTIONS,
    compose_utterance_instruct,
)
from umbraaudio.voices.qwen_tts import QwenTTSEngine


class _FakeResponse:
    def __init__(self, content: bytes) -> None:
        self.status_code = 200
        self.content = content

    @property
    def headers(self) -> dict[str, str]:
        return {"content-type": "audio/wav"}

    def json(self) -> dict[str, object]:
        return {}


def _tiny_wav_bytes() -> bytes:
    sr = 24000
    t = np.linspace(0, 0.1, int(sr * 0.1), dtype=np.float32)
    data = (0.1 * np.sin(2 * np.pi * 220 * t)).astype(np.float32)
    buf = io.BytesIO()
    sf.write(buf, data, sr, format="WAV")
    return buf.getvalue()


class TestComposeUtteranceInstruct:
    def test_no_variation_returns_none(self) -> None:
        assert compose_utterance_instruct() is None
        assert compose_utterance_instruct(emotion="neutral", pace="moderate") is None
        assert compose_utterance_instruct("", "neutral", "moderate") is None

    def test_custom_instruct_only(self) -> None:
        assert compose_utterance_instruct("Growl this line.") == "Growl this line."

    def test_emotion_adds_prescriptive_text(self) -> None:
        result = compose_utterance_instruct(emotion="angry")
        assert "anger" in result
        assert result == EMOTION_INSTRUCT_MAP["angry"]

    def test_neutral_emotion_contributes_nothing(self) -> None:
        result = compose_utterance_instruct("Say it plain.", emotion="neutral")
        assert result == "Say it plain."

    def test_pace_drives_tempo(self) -> None:
        fast = compose_utterance_instruct(pace="fast")
        assert fast and "quick" in fast
        slow = compose_utterance_instruct(pace="slow")
        assert slow and "slow" in slow
        assert compose_utterance_instruct(pace="moderate") is None

    def test_unknown_emotion_is_skipped(self) -> None:
        result = compose_utterance_instruct("Keep this.", emotion="not-a-real-emotion")
        assert result == "Keep this."

    def test_combines_custom_emotion_pace(self) -> None:
        result = compose_utterance_instruct(
            "Hold the line.", emotion="furious", pace="fast"
        )
        assert result.startswith("Hold the line.")
        assert "rage" in result
        assert "quick" in result
        assert len(result.split(". ")) == 1 + 2 + 1

    def test_every_emotion_option_has_a_directive_except_neutral(self) -> None:
        for option in EMOTION_OPTIONS:
            assert option in EMOTION_INSTRUCT_MAP, f"missing {option} in map"
            if option == "neutral":
                assert EMOTION_INSTRUCT_MAP[option] == ""
            else:
                assert EMOTION_INSTRUCT_MAP[option].strip(), (
                    f"{option} has empty directive"
                )

    def test_map_keys_stay_in_sync_with_options(self) -> None:
        assert set(EMOTION_INSTRUCT_MAP) == set(EMOTION_OPTIONS)

    def test_pace_options_stay_in_sync_with_map(self) -> None:
        assert set(PACE_INSTRUCT) == set(PACE_OPTIONS)
        assert PACE_INSTRUCT["moderate"] == ""
        assert PACE_INSTRUCT["slow"] and PACE_INSTRUCT["fast"]


class TestNarratorProfile:
    def test_narrator_compression_targets_calm_even_delivery(self) -> None:
        assert "Do not dramatize" in NARRATOR_COMPRESSION_INSTRUCT
        assert "steadily" in NARRATOR_COMPRESSION_INSTRUCT


class TestRemoteInstructComposition:
    def test_instruct_sent_alongside_voice_ref(self, tmp_path, monkeypatch) -> None:
        ref = tmp_path / "narrator_ref.wav"
        ref.write_bytes(_tiny_wav_bytes())
        captured: dict[str, object] = {}

        def fake_post(
            url: str,
            data: dict[str, object],
            files: dict[str, object],
            timeout: float,
        ) -> _FakeResponse:
            captured["data"] = data
            captured["files"] = files
            return _FakeResponse(_tiny_wav_bytes())

        monkeypatch.setattr(httpx, "post", fake_post)

        engine = QwenTTSEngine(qwen_tts_url="http://tts.test:8881")
        out = tmp_path / "out.wav"
        engine.generate(
            text="Hello.",
            output_path=out,
            voice_ref=ref,
            instruct=EMOTION_INSTRUCT_MAP["angry"],
        )

        data: dict[str, object] = captured["data"]
        assert data["voice_description"] == EMOTION_INSTRUCT_MAP["angry"]
        assert captured["files"], "ref audio must still be attached"
        assert out.exists()

    def test_no_instruct_means_no_voice_description(self, tmp_path, monkeypatch) -> None:
        ref = tmp_path / "kainda_ref.wav"
        ref.write_bytes(_tiny_wav_bytes())
        captured: dict[str, object] = {}

        def fake_post(
            url: str,
            data: dict[str, object],
            files: dict[str, object],
            timeout: float,
        ) -> _FakeResponse:
            captured["data"] = data
            return _FakeResponse(_tiny_wav_bytes())

        monkeypatch.setattr(httpx, "post", fake_post)

        engine = QwenTTSEngine(qwen_tts_url="http://tts.test:8881")
        engine.generate(text="Hello.", output_path=tmp_path / "out.wav", voice_ref=ref)

        assert "voice_description" not in captured["data"]