xiaoxia-saas/tests/unit/test_subtitle_generator.py

"""字幕生成器 + Mock ASR 单元测试。"""

import tempfile
from pathlib import Path

import pytest

from apps.worker.video_processing.subtitle_generator import (
    _wrap_text,
    generate_ass_from_timeline,
)
from packages.adapters.asr.mock_asr_service import MockASRService
from packages.domain.subtitle import SubtitleSegment, SubtitleTimeline
from packages.ports.asr_service import ASRServiceError


class TestMockASRService:
    def test_transcribe_with_mock_text(self):
        service = MockASRService(mock_text="你好世界！这是一段测试语音识别的文字。用来验证Mock ASR是否正常工作。")

        # 创建一个假的音频文件（mock不真的读内容）
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            f.write(b"fake audio data")
            audio_path = Path(f.name)

        try:
            timeline = service.transcribe(audio_path, language="zh")
            assert timeline is not None
            assert timeline.language == "zh"
            assert timeline.segment_count > 0
            assert timeline.total_duration > 0
            # 总字数应该对得上
            assert timeline.total_chars == len("你好世界！这是一段测试语音识别的文字。用来验证Mock ASR是否正常工作。")
        finally:
            audio_path.unlink()

    def test_transcribe_file_not_found(self):
        service = MockASRService()
        with pytest.raises(ASRServiceError):
            service.transcribe(Path("/nonexistent/audio.wav"))

    def test_transcribe_with_word_timestamps(self):
        service = MockASRService(mock_text="你好世界！")

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            f.write(b"fake")
            audio_path = Path(f.name)

        try:
            timeline = service.transcribe(audio_path, with_word_timestamps=True)
            # 每段应该有词级时间戳
            for seg in timeline.segments:
                if seg.words:
                    assert len(seg.words) > 0
                    assert seg.words[0].start >= seg.start
                    assert seg.words[-1].end <= seg.end
        finally:
            audio_path.unlink()

    def test_auto_detect_language(self):
        service = MockASRService(mock_text="Hello world. This is a test.")

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
            f.write(b"fake")
            audio_path = Path(f.name)

        try:
            timeline = service.transcribe(audio_path, language=None)
            # None 时默认 zh
            assert timeline.language == "zh"
        finally:
            audio_path.unlink()


class TestWrapText:
    def test_short_text_no_wrap(self):
        result = _wrap_text("你好世界", 20)
        assert result == ["你好世界"]

    def test_wrap_at_punctuation(self):
        result = _wrap_text("你好世界！这是一段很长的测试文字。", 10)
        assert len(result) == 2
        assert "！" in result[0]

    def test_hard_wrap_no_punctuation(self):
        result = _wrap_text("一二三四五六七八九十一二三四五六七八九十", 10)
        assert len(result) == 2
        assert len(result[0]) == 10
        assert len(result[1]) == 10

    def test_exact_length(self):
        result = _wrap_text("一二三四五六七八九十", 10)
        assert len(result) == 1


class TestGenerateAssFromTimeline:
    def test_generate_basic(self):
        timeline = SubtitleTimeline(
            segments=[
                SubtitleSegment(text="你好世界", start=0.0, end=2.0),
                SubtitleSegment(text="这是测试", start=2.0, end=4.0),
            ],
            language="zh",
            total_duration=4.0,
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "test.ass"
            result = generate_ass_from_timeline(
                output_path,
                timeline,
                video_width=1920,
                video_height=1080,
            )

            assert result.exists()
            content = result.read_text(encoding="utf-8")
            assert "[Script Info]" in content
            assert "[V4+ Styles]" in content
            assert "[Events]" in content
            assert "你好世界" in content
            assert "这是测试" in content
            assert "PlayResX: 1920" in content
            assert "PlayResY: 1080" in content

    def test_empty_timeline(self):
        timeline = SubtitleTimeline(segments=[])

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "empty.ass"
            result = generate_ass_from_timeline(output_path, timeline, video_width=1920, video_height=1080)
            assert result.exists()
            assert result.read_text(encoding="utf-8") == ""

    def test_with_custom_style(self):
        timeline = SubtitleTimeline(
            segments=[SubtitleSegment(text="测试", start=0, end=1)],
            total_duration=1.0,
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "style.ass"
            generate_ass_from_timeline(
                output_path,
                timeline,
                video_width=1280,
                video_height=720,
                subtitle_config={
                    "font": "微软雅黑",
                    "size": 32,
                    "color": "#ff0000",
                    "position": "bottom",
                },
            )

            content = output_path.read_text(encoding="utf-8")
            assert "微软雅黑" in content
            assert "32" in content

    def test_time_format(self):
        timeline = SubtitleTimeline(
            segments=[
                SubtitleSegment(text="测试", start=0.5, end=1.25),
                SubtitleSegment(text="长字幕", start=3661.0, end=3662.5),  # 超过1小时
            ],
            total_duration=3662.5,
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "time.ass"
            generate_ass_from_timeline(output_path, timeline, video_width=1920, video_height=1080)

            content = output_path.read_text(encoding="utf-8")
            # 0:00:00.50 格式
            assert "0:00:00.50" in content
            assert "0:00:01.25" in content
            # 1:01:01.00 格式（3661秒 = 1小时1分1秒）
            assert "1:01:01.00" in content