tests/test_chat_tts_stream.py

import pytest

from app.services.chat_tts_stream import ChatTTSStream


class RecordingManager:
    def __init__(self) -> None:
        self.begin_calls: list[dict] = []
        self.enqueued: list[tuple[str, int, str]] = []
        self.finished: list[str] = []

    async def begin(self, project_id, message_id, send, **kwargs):
        self.begin_calls.append(kwargs)
        return "generation"

    async def enqueue(self, generation_id, sentence_index, text):
        self.enqueued.append((generation_id, sentence_index, text))

    async def finish(self, generation_id):
        self.finished.append(generation_id)


@pytest.mark.asyncio
async def test_voice_originated_stream_buffers_one_complete_native_utterance():
    manager = RecordingManager()
    stream = await ChatTTSStream.create(
        manager=manager,
        project_id="project",
        message_id="message",
        send=lambda *_args: None,
        input_mode="voice",
        auto_read_reply=True,
    )

    await stream.feed("First sentence. Sec")
    await stream.feed("ond sentence!")
    assert manager.enqueued == []
    await stream.finish()

    assert manager.enqueued == [("generation", 0, "First sentence. Second sentence!")]
    assert manager.finished == ["generation"]


@pytest.mark.asyncio
async def test_auto_read_forwards_the_model_and_voice_selection():
    manager = RecordingManager()

    await ChatTTSStream.create(
        manager=manager,
        project_id="project",
        message_id="message",
        send=lambda *_args: None,
        input_mode="voice",
        auto_read_reply=True,
        model_id="piper",
        voice_id="en_US-lessac-medium",
    )

    assert manager.begin_calls == [{"model_id": "piper", "voice_id": "en_US-lessac-medium"}]


@pytest.mark.asyncio
@pytest.mark.parametrize(("input_mode", "auto_read_reply"), [("text", True), ("voice", False)])
async def test_stream_is_disabled_for_text_or_opted_out_turns(input_mode, auto_read_reply):
    manager = RecordingManager()
    stream = await ChatTTSStream.create(
        manager=manager,
        project_id="project",
        message_id="message",
        send=lambda *_args: None,
        input_mode=input_mode,
        auto_read_reply=auto_read_reply,
    )
    await stream.feed("Nothing should play.")
    await stream.finish()

    assert stream.generation_id is None
    assert manager.begin_calls == []