unios-tts/tests/test_synthesis.py

import numpy as np
import torch

from app.segmentation import Segment
from app.synthesis import SynthesisEngine, ModelEntry


class FakeModel:
    def __init__(self, name, calls):
        self.name = name
        self.calls = calls
        self.speakers = ["spk1", "spk2"]

    def apply_tts(self, *, text=None, ssml_text=None, speaker, sample_rate,
                  put_accent, put_yo):
        self.calls.append({
            "name": self.name, "speaker": speaker, "text": text,
            "ssml_text": ssml_text, "put_accent": put_accent,
        })
        return torch.zeros(10)


class FakeTts:
    def __init__(self, name, calls, put_accent=True):
        self.sample_rate = 24000
        self.put_accent = put_accent
        self.put_yo = True
        self.tts_model = FakeModel(name, calls)


def make_engine(use_ssml=True, put_accent=True):
    calls: list[dict] = []
    engine = SynthesisEngine(
        models=[
            ModelEntry(model_id="ru_model", tts=FakeTts("ru", calls, put_accent), languages=["ru"]),
            ModelEntry(model_id="en_model", tts=FakeTts("en", calls, put_accent), languages=["en"]),
        ],
        sample_rate=24000,
        use_ssml=use_ssml,
    )
    return engine, calls


def test_single_speaker_for_all_segments():
    engine, calls = make_engine(use_ssml=False)
    segments = [Segment(text="Озон", lang="ru"), Segment(text="Пайтон", lang="ru")]
    pcm = engine.synthesize_segments(segments, "baya", "ru")
    assert pcm
    assert len(calls) == 2
    assert all(c["name"] == "ru" and c["speaker"] == "baya" for c in calls)


def test_en_lang_selects_en_model():
    engine, calls = make_engine(use_ssml=False)
    segments = [Segment(text="hello", lang="en")]
    engine.synthesize_segments(segments, "en_0", "en")
    assert calls[0]["name"] == "en"
    assert calls[0]["text"] == "hello"


def test_legacy_path_uses_text_not_ssml():
    engine, calls = make_engine(use_ssml=False)
    engine.synthesize_segments([Segment(text="привет", lang="ru")], "baya", "ru")
    assert calls[0]["text"] == "привет"
    assert calls[0]["ssml_text"] is None


def test_no_silence_padding_between_segments():
    engine, calls = make_engine(use_ssml=False)
    segments = [Segment(text="а", lang="ru"), Segment(text="б", lang="ru")]
    pcm = engine.synthesize_segments(segments, "baya", "ru")
    expected = (np.zeros(10) * 32767).astype(np.int16)
    assert len(pcm) == 2 * len(expected.tobytes())


# --- ssml_text path (task 3) ---

def test_ssml_path_passes_ssml_text():
    engine, calls = make_engine()
    segments = [Segment(text="привет", lang="ru"), Segment(text="мир", lang="ru")]
    pcm = engine.synthesize_segments(segments, "baya", "ru")
    assert pcm
    assert len(calls) == 1  # один чанк на абзац
    c = calls[0]
    assert c["text"] is None
    assert c["ssml_text"].startswith("<speak>")
    assert "<s>привет</s>" in c["ssml_text"]
    assert "<s>мир</s>" in c["ssml_text"]


def test_ssml_path_single_call_per_chunk():
    engine, calls = make_engine()
    # Два разных абзаца -> два <p> -> (в пределах лимита) один документ/один вызов.
    segments = [
        Segment(text="первый", lang="ru", para=0),
        Segment(text="второй", lang="ru", para=1),
    ]
    engine.synthesize_segments(segments, "baya", "ru")
    assert len(calls) == 1
    assert "<p>" in calls[0]["ssml_text"]


def test_ssml_prosody_rate_and_pitch_passed_to_builder():
    engine, calls = make_engine()
    engine.synthesize_segments(
        [Segment(text="привет", lang="ru")], "baya", "ru",
        rate="slow", pitch="high",
    )
    ssml = calls[0]["ssml_text"]
    assert '<prosody' in ssml
    assert 'rate="slow"' in ssml
    assert 'pitch="high"' in ssml


def test_ssml_medium_prosody_omitted():
    engine, calls = make_engine()
    engine.synthesize_segments(
        [Segment(text="привет", lang="ru")], "baya", "ru",
        rate="medium", pitch="medium",
    )
    assert "<prosody" not in calls[0]["ssml_text"]


def test_ssml_use_ssml_false_overrides_for_call():
    engine, calls = make_engine(use_ssml=True)
    engine.synthesize_segments(
        [Segment(text="привет", lang="ru")], "baya", "ru", use_ssml=False
    )
    assert calls[0]["ssml_text"] is None
    assert calls[0]["text"] == "привет"


def test_put_accent_true_on_ssml_path():
    # put_accent=True даже если модель по умолчанию False (ручные + работают).
    engine, calls = make_engine(put_accent=False)
    engine.synthesize_segments([Segment(text="привет", lang="ru")], "baya", "ru")
    assert calls[0]["put_accent"] is True


def test_manual_stress_marker_reaches_model_via_ssml():
    engine, calls = make_engine()
    engine.synthesize_segments(
        [Segment(text="втор+ой сл+ово", lang="ru")], "baya", "ru"
    )
    assert "втор+ой" in calls[0]["ssml_text"]
    assert "сл+ово" in calls[0]["ssml_text"]


def test_manual_stress_non_vowel_passes_without_error():
    engine, calls = make_engine()
    pcm = engine.synthesize_segments(
        [Segment(text="т+ест", lang="ru")], "baya", "ru"
    )
    assert pcm
    assert "т+ест" in calls[0]["ssml_text"]