xiaoxia-saas/tests/unit/test_tts_streaming.py

"""P2 WebSocket 流式合成单元测试。

覆盖：
- TTSStreamingService 流式合成逻辑
- 短文本流式合成
- 长文本分段流式合成
- 错误处理
"""

from __future__ import annotations

from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from packages.application.cosyvoice_service import CosyVoiceError, CosyVoiceService
from packages.application.tts_job.streaming_service import (
    TTSStreamingService,
)


class MockWebSocket:
    """Mock WebSocket for testing."""

    def __init__(self):
        self.sent_json = []
        self.sent_bytes = []
        self.accepted = False

    async def accept(self):
        self.accepted = True

    async def send_json(self, data):
        self.sent_json.append(data)

    async def send_bytes(self, data):
        self.sent_bytes.append(data)

    async def receive_json(self):
        return {}


class TestTTSStreamingService:
    """测试 TTS 流式合成服务。"""

    @pytest.mark.asyncio
    async def test_empty_text_returns_error(self):
        """空文本返回错误。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        params = {"text": "", "voice_id": "test_voice"}

        await service.synthesize_and_stream(ws, params)

        assert len(ws.sent_json) == 1
        assert ws.sent_json[0]["type"] == "error"
        assert "文本不能为空" in ws.sent_json[0]["message"]

    @pytest.mark.asyncio
    async def test_text_too_long_returns_error(self):
        """超长文本返回错误。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        params = {"text": "x" * 10001, "voice_id": "test_voice"}

        await service.synthesize_and_stream(ws, params)

        assert len(ws.sent_json) == 1
        assert ws.sent_json[0]["type"] == "error"
        assert "文本过长" in ws.sent_json[0]["message"]

    @pytest.mark.asyncio
    async def test_short_text_stream_success(self):
        """短文本流式合成成功。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "task_id": "task_1",
            "audio_url": "https://temp.com/audio.mp3",
            "duration": 5.0,
            "file_size": 10000,
        }

        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        params = {"text": "测试文本", "voice_id": "test_voice", "format": "mp3"}

        with patch.object(service, "_download_audio", return_value=b"fake audio data"):
            await service.synthesize_and_stream(ws, params)

        # 验证发送了 started 帧
        assert ws.sent_json[0]["type"] == "started"
        assert ws.sent_json[0]["segment_count"] == 1

        # 验证发送了二进制音频数据
        assert len(ws.sent_bytes) > 0

        # 验证发送了 done 帧
        assert ws.sent_json[-1]["type"] == "done"
        assert ws.sent_json[-1]["format"] == "mp3"

    @pytest.mark.asyncio
    async def test_short_text_cosyvoice_error(self):
        """短文本合成时 CosyVoice 报错。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.side_effect = CosyVoiceError("API error")

        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        params = {"text": "测试文本", "voice_id": "test_voice"}

        await service.synthesize_and_stream(ws, params)

        # 验证发送了 started 帧和 error 帧
        assert len(ws.sent_json) == 2
        assert ws.sent_json[0]["type"] == "started"
        assert ws.sent_json[1]["type"] == "error"
        assert "API error" in ws.sent_json[1]["message"]

    @pytest.mark.asyncio
    async def test_short_text_no_audio_url(self):
        """短文本合成未返回 audio_url。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "task_id": "task_1",
            "audio_url": "",
        }

        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        params = {"text": "测试文本", "voice_id": "test_voice"}

        await service.synthesize_and_stream(ws, params)

        # 验证发送了 started 帧和 error 帧
        assert len(ws.sent_json) == 2
        assert ws.sent_json[0]["type"] == "started"
        assert ws.sent_json[1]["type"] == "error"
        assert "未返回音频 URL" in ws.sent_json[1]["message"]

    @pytest.mark.asyncio
    async def test_long_text_stream_success(self):
        """长文本分段流式合成成功。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        # 每个分段都返回 audio_url
        cosyvoice.submit_synthesize_task.side_effect = [
            {"task_id": "", "audio_url": "https://temp.com/seg1.mp3", "duration": 2.0},
            {"task_id": "", "audio_url": "https://temp.com/seg2.mp3", "duration": 3.0},
        ]

        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        # 超过 500 字的文本
        params = {"text": "x" * 600, "voice_id": "test_voice", "format": "mp3"}

        with patch.object(service, "_download_audio", return_value=b"segment audio"):
            await service.synthesize_and_stream(ws, params)

        # 验证发送了 started 帧（多段）
        assert ws.sent_json[0]["type"] == "started"
        assert ws.sent_json[0]["segment_count"] >= 2

        # 验证发送了 segment_done 帧
        segment_done_count = sum(1 for msg in ws.sent_json if msg["type"] == "segment_done")
        assert segment_done_count >= 2

        # 验证发送了 done 帧
        assert ws.sent_json[-1]["type"] == "done"

    @pytest.mark.asyncio
    async def test_long_text_segment_failure(self):
        """长文本分段合成失败。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.side_effect = CosyVoiceError("Segment error")

        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        params = {"text": "x" * 600, "voice_id": "test_voice"}

        await service.synthesize_and_stream(ws, params)

        # 验证发送了错误帧
        error_msgs = [msg for msg in ws.sent_json if msg["type"] == "error"]
        assert len(error_msgs) > 0
        assert "合成失败" in error_msgs[0]["message"]

    @pytest.mark.asyncio
    async def test_stream_audio_chunks(self):
        """音频分块推送。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        audio_data = b"x" * 10000  # 10KB

        total = await service._stream_audio_chunks(ws, audio_data)

        assert total == 10000
        # 验证分块发送（4KB per chunk）
        assert len(ws.sent_bytes) == 3  # 4096 + 4096 + 1808

    @pytest.mark.asyncio
    async def test_send_json_suppresses_exceptions(self):
        """_send_json 抑制异常。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)

        ws = MockWebSocket()
        ws.send_json = AsyncMock(side_effect=Exception("Send failed"))

        # 不应该抛出异常
        await service._send_json(ws, {"type": "test"})

    @pytest.mark.asyncio
    async def test_download_audio(self):
        """下载音频数据（SSRF防护走safe_download_bytes，mock掉安全层）。"""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)

        with patch("packages.application.tts_job.streaming_service.safe_download_bytes") as mock_download:
            mock_download.return_value = b"audio data"

            result = service._download_audio("https://example.com/audio.mp3")

            assert result == b"audio data"
            mock_download.assert_called_once()


class TestTTSStreamingEdgeCases:
    """流式合成边界测试."""

    @pytest.mark.asyncio
    async def test_exactly_500_chars_uses_short_text(self):
        """刚好500字走短文本路径."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/audio.mp3",
            "duration": 5.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "x" * 500, "voice_id": "test_voice"}

        with patch.object(service, "_download_audio", return_value=b"audio"):
            await service.synthesize_and_stream(ws, params)

        # 短文本只有1个segment
        assert ws.sent_json[0]["type"] == "started"
        assert ws.sent_json[0]["segment_count"] == 1

    @pytest.mark.asyncio
    async def test_501_chars_uses_long_text(self):
        """501字走长文本分段路径."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/audio.mp3",
            "duration": 2.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "x" * 501, "voice_id": "test_voice"}

        with patch.object(service, "_download_audio", return_value=b"audio"):
            await service.synthesize_and_stream(ws, params)

        # 长文本segment_count > 1
        assert ws.sent_json[0]["type"] == "started"
        assert ws.sent_json[0]["segment_count"] >= 2

    @pytest.mark.asyncio
    async def test_short_text_speed_param_passed(self):
        """短文本合成时速度参数正确传递."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/audio.mp3",
            "duration": 3.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "测试", "voice_id": "v1", "speed": 1.5, "format": "wav"}

        with patch.object(service, "_download_audio", return_value=b"audio"):
            await service.synthesize_and_stream(ws, params)

        cosyvoice.submit_synthesize_task.assert_called_once()
        call_kwargs = cosyvoice.submit_synthesize_task.call_args.kwargs
        assert call_kwargs["speed"] == 1.5
        assert call_kwargs["format"] == "wav"
        assert call_kwargs["voice_id"] == "v1"

    @pytest.mark.asyncio
    async def test_short_text_sample_rate_param(self):
        """短文本合成时采样率参数传递."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/audio.mp3",
            "duration": 1.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "测试", "voice_id": "v1", "sample_rate": 44100}

        with patch.object(service, "_download_audio", return_value=b"audio"):
            await service.synthesize_and_stream(ws, params)

        call_kwargs = cosyvoice.submit_synthesize_task.call_args.kwargs
        assert call_kwargs["sample_rate"] == 44100

    @pytest.mark.asyncio
    async def test_single_chunk_audio(self):
        """小于4KB的音频只发1块."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        audio_data = b"x" * 1000  # 1KB < 4KB

        total = await service._stream_audio_chunks(ws, audio_data)
        assert total == 1000
        assert len(ws.sent_bytes) == 1
        assert ws.sent_bytes[0] == audio_data

    @pytest.mark.asyncio
    async def test_exact_chunk_size_audio(self):
        """刚好4KB的音频只发1块."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        audio_data = b"x" * 4096

        total = await service._stream_audio_chunks(ws, audio_data)
        assert total == 4096
        assert len(ws.sent_bytes) == 1

    @pytest.mark.asyncio
    async def test_empty_audio_chunks(self):
        """空音频数据不发送任何块."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()

        total = await service._stream_audio_chunks(ws, b"")
        assert total == 0
        assert len(ws.sent_bytes) == 0

    @pytest.mark.asyncio
    async def test_short_text_unexpected_exception(self):
        """短文本合成时非预期异常捕获."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.side_effect = RuntimeError("Unexpected error")

        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "测试", "voice_id": "v1"}

        await service.synthesize_and_stream(ws, params)

        assert ws.sent_json[-1]["type"] == "error"
        assert "合成失败" in ws.sent_json[-1]["message"]

    @pytest.mark.asyncio
    async def test_short_text_download_failure(self):
        """短文本音频下载失败."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/audio.mp3",
            "duration": 1.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "测试", "voice_id": "v1"}

        with patch.object(service, "_download_audio", side_effect=Exception("Download failed")):
            await service.synthesize_and_stream(ws, params)

        assert ws.sent_json[-1]["type"] == "error"
        assert "音频推送失败" in ws.sent_json[-1]["message"]

    @pytest.mark.asyncio
    async def test_long_text_segment_count_matches_split(self):
        """长文本分段数量与split_text结果一致."""
        from packages.application.tts_job.text_splitter import split_text

        text = "x" * 1200
        segments = split_text(text, max_chars=500)
        expected_count = len(segments)

        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/a.mp3",
            "duration": 1.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": text, "voice_id": "v1"}

        with patch.object(service, "_download_audio", return_value=b"audio"):
            await service.synthesize_and_stream(ws, params)

        assert ws.sent_json[0]["segment_count"] == expected_count
        segment_done = sum(1 for m in ws.sent_json if m["type"] == "segment_done")
        assert segment_done == expected_count

    @pytest.mark.asyncio
    async def test_long_text_total_bytes_accumulated(self):
        """长文本总字节数正确累加."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        cosyvoice.submit_synthesize_task.return_value = {
            "audio_url": "https://temp.com/a.mp3",
            "duration": 1.0,
        }
        service = TTSStreamingService(cosyvoice)
        ws = MockWebSocket()
        params = {"text": "x" * 600, "voice_id": "v1"}

        audio_chunk = b"x" * 5000
        with patch.object(service, "_download_audio", return_value=audio_chunk):
            await service.synthesize_and_stream(ws, params)

        # done帧中file_size应为分段数 * 每段大小
        done_msg = ws.sent_json[-1]
        assert done_msg["type"] == "done"
        segment_count = ws.sent_json[0]["segment_count"]
        assert done_msg["file_size"] == segment_count * 5000

    def test_download_audio_passes_purpose_and_mime(self):
        """_download_audio正确传递参数给safe_download_bytes."""
        cosyvoice = MagicMock(spec=CosyVoiceService)
        service = TTSStreamingService(cosyvoice)

        with patch("packages.application.tts_job.streaming_service.safe_download_bytes") as mock:
            mock.return_value = b"data"
            service._download_audio("https://example.com/a.wav")

            mock.assert_called_once()
            kwargs = mock.call_args.kwargs
            assert kwargs["purpose"] == "tts_streaming_download"
            assert kwargs["timeout"] == 60.0
            assert "allowed_mime_types" in kwargs