xiaoxia-saas/tests/unit/test_audio_merger.py

"""音频合并器单元测试."""

from __future__ import annotations

import os
import tempfile
from unittest.mock import MagicMock, patch

import pytest

from packages.application.tts_job.audio_merger import AudioMergeError, AudioMerger


@pytest.fixture
def sample_audio_dir():
    """创建临时目录，放几个模拟音频文件"""
    tmpdir = tempfile.mkdtemp()
    files = []
    for i in range(3):
        fpath = os.path.join(tmpdir, f"part{i}.mp3")
        with open(fpath, "wb") as f:
            f.write(f"audio_data_{i}".encode() * 100)
        files.append(fpath)
    yield files
    import shutil

    shutil.rmtree(tmpdir, ignore_errors=True)


class TestAudioMerger:
    """AudioMerger 测试"""

    def test_empty_list_raises_error(self):
        """空列表抛出 AudioMergeError"""
        merger = AudioMerger()
        with pytest.raises(AudioMergeError, match="没有可合并的音频文件"):
            merger.merge([])

    def test_single_file_returns_content(self, sample_audio_dir):
        """单文件直接返回文件内容"""
        merger = AudioMerger()
        result = merger.merge([sample_audio_dir[0]])

        with open(sample_audio_dir[0], "rb") as f:
            expected = f.read()

        assert result == expected

    def test_single_file_no_ffmpeg_needed(self, sample_audio_dir):
        """单文件不需要调用 FFmpeg"""
        with patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg:
            merger = AudioMerger()
            merger.merge([sample_audio_dir[0]])
            mock_ffmpeg.assert_not_called()

    def test_merge_multiple_files(self, sample_audio_dir):
        """多文件合并调用 FFmpeg"""
        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
        ):
            # 模拟 FFmpeg 成功：在 output_path 写点数据
            def fake_run_ffmpeg(cmd, timeout=120):
                output_idx = cmd.index("-c") + 2  # -c copy 后面是 output_path
                output_path = cmd[-1]
                with open(output_path, "wb") as f:
                    f.write(b"merged_audio_data")
                return MagicMock(stdout=b"", stderr=b"")

            mock_ffmpeg.side_effect = fake_run_ffmpeg

            merger = AudioMerger()
            result = merger.merge(sample_audio_dir)

            assert result == b"merged_audio_data"
            mock_ffmpeg.assert_called_once()

    def test_merge_concat_list_generated(self, sample_audio_dir):
        """生成正确的 concat demuxer 列表文件"""
        import subprocess

        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
        ):

            captured_list_content = []

            def fake_run_ffmpeg(cmd, timeout=120):
                # 找到 -i 参数后面的文件路径
                # 命令结构: ffmpeg -y -f concat -safe 0 -i LIST_PATH -c copy OUTPUT
                for i, arg in enumerate(cmd):
                    if arg == "-i" and i + 1 < len(cmd):
                        list_path = cmd[i + 1]
                        if list_path.endswith(".txt"):
                            with open(list_path, "r") as f:
                                captured_list_content.append(f.read())
                        break
                # 写输出文件
                output_path = cmd[-1]
                with open(output_path, "wb") as f:
                    f.write(b"fake")
                return MagicMock(stdout=b"", stderr=b"")

            mock_ffmpeg.side_effect = fake_run_ffmpeg

            merger = AudioMerger()
            merger.merge(sample_audio_dir, output_format="mp3")

            # 检查列表文件包含所有输入文件
            assert len(captured_list_content) == 1
            list_content = captured_list_content[0]
            for fpath in sample_audio_dir:
                assert fpath in list_content.replace("'\\''", "'")

    def test_merge_ffmpeg_failure_raises(self, sample_audio_dir):
        """FFmpeg 失败抛出 AudioMergeError"""
        from subprocess import CalledProcessError

        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
        ):
            mock_ffmpeg.side_effect = CalledProcessError(returncode=1, cmd=["ffmpeg"], stderr=b"error message")

            merger = AudioMerger()
            with pytest.raises(AudioMergeError, match="FFmpeg 合并失败"):
                merger.merge(sample_audio_dir)

    def test_merge_timeout_raises(self, sample_audio_dir):
        """合并超时抛出 AudioMergeError"""
        from subprocess import TimeoutExpired

        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
        ):
            mock_ffmpeg.side_effect = TimeoutExpired(cmd=["ffmpeg"], timeout=120)

            merger = AudioMerger()
            with pytest.raises(AudioMergeError, match="超时"):
                merger.merge(sample_audio_dir)

    def test_merge_cleanup_temp_dir(self, sample_audio_dir):
        """合并完成后清理临时目录"""
        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
            patch("packages.application.tts_job.audio_merger.shutil.rmtree") as mock_rmtree,
        ):

            def fake_run_ffmpeg(cmd, timeout=120):
                output_path = cmd[-1]
                with open(output_path, "wb") as f:
                    f.write(b"data")
                return MagicMock()

            mock_ffmpeg.side_effect = fake_run_ffmpeg

            merger = AudioMerger()
            merger.merge(sample_audio_dir)

            mock_rmtree.assert_called_once()
            # 第一个参数是临时目录路径
            temp_dir_path = mock_rmtree.call_args[0][0]
            assert "tts_merge_" in temp_dir_path

    def test_merge_cleanup_on_error(self, sample_audio_dir):
        """合并失败也清理临时目录"""
        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
            patch("packages.application.tts_job.audio_merger.shutil.rmtree") as mock_rmtree,
        ):
            from subprocess import CalledProcessError

            mock_ffmpeg.side_effect = CalledProcessError(1, ["ffmpeg"])

            merger = AudioMerger()
            try:
                merger.merge(sample_audio_dir)
            except AudioMergeError:
                pass

            mock_rmtree.assert_called_once()

    def test_merge_custom_output_format(self, sample_audio_dir):
        """自定义输出格式"""
        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
        ):

            def fake_run_ffmpeg(cmd, timeout=120):
                output_path = cmd[-1]
                assert output_path.endswith(".wav")
                with open(output_path, "wb") as f:
                    f.write(b"data")
                return MagicMock()

            mock_ffmpeg.side_effect = fake_run_ffmpeg

            merger = AudioMerger()
            merger.merge(sample_audio_dir, output_format="wav")

    def test_merge_two_files(self, sample_audio_dir):
        """两个文件合并"""
        with (
            patch("packages.application.tts_job.audio_merger.run_ffmpeg") as mock_ffmpeg,
            patch("packages.application.tts_job.audio_merger.FFMPEG_BIN", "ffmpeg"),
        ):

            def fake_run_ffmpeg(cmd, timeout=120):
                output_path = cmd[-1]
                with open(output_path, "wb") as f:
                    f.write(b"two_files_merged")
                return MagicMock()

            mock_ffmpeg.side_effect = fake_run_ffmpeg

            merger = AudioMerger()
            result = merger.merge(sample_audio_dir[:2])
            assert result == b"two_files_merged"