hunjing/yuanqi/tests/test_asr.py

"""
ASR 引擎和端点测试。
"""

from __future__ import annotations

import io
import wave
from unittest.mock import patch

import pytest

from src.endpoints.asr import _ensure_loaded, _validate_format
from src.engine.asr import _SUPPORTED_FORMATS, SenseVoiceASREngine


def _make_test_wav(duration_sec: float = 0.5, sample_rate: int = 16000) -> bytes:
    """生成测试用的 WAV 文件字节（静音音频）。"""
    buf = io.BytesIO()
    num_samples = int(duration_sec * sample_rate)
    with wave.open(buf, "wb") as wf:
        wf.setnchannels(1)
        wf.setsampwidth(2)
        wf.setframerate(sample_rate)
        wf.writeframes(b"\x00\x00" * num_samples)
    return buf.getvalue()


# ── 测试 _SUPPORTED_FORMATS ────────────────────────────────────────────────


class TestSupportedFormats:
    def test_common_formats_present(self) -> None:
        """应包含常见音频格式。"""
        for fmt in ("wav", "mp3", "m4a", "flac", "ogg", "opus", "aac"):
            assert fmt in _SUPPORTED_FORMATS, f"缺少常用格式: {fmt}"


# ── 测试 _validate_format ──────────────────────────────────────────────────


class TestValidateFormat:
    def test_supported_format(self) -> None:
        """支持的格式返回扩展名。"""
        assert _validate_format("audio.wav") == "wav"
        assert _validate_format("speech.mp3") == "mp3"

    @pytest.mark.asyncio
    async def test_no_extension_raises(self) -> None:
        """无扩展名应抛出异常。"""
        from fastapi import HTTPException

        with pytest.raises(HTTPException) as exc:
            _validate_format("no_ext")
        assert exc.value.status_code == 400
        assert str(exc.value.detail["code"]) == "unsupported_format"

    @pytest.mark.asyncio
    async def test_unsupported_format(self) -> None:
        """不支持的格式应抛出异常。"""
        from fastapi import HTTPException

        with pytest.raises(HTTPException) as exc:
            _validate_format("video.avi")
        assert exc.value.status_code == 400
        assert str(exc.value.detail["code"]) == "unsupported_format"


# ── 测试 SenseVoiceASREngine ──────────────────────────────────────────────


class TestSenseVoiceASREngine:
    def test_init(self) -> None:
        """引擎初始化应使用默认配置。"""
        engine = SenseVoiceASREngine()
        assert engine.name == "asr"
        assert engine.model_loaded is False
        assert engine._model is None

    @pytest.mark.asyncio
    async def test_health_unloaded(self) -> None:
        """模型未加载时 health 应返回 unavailable。"""
        engine = SenseVoiceASREngine()
        h = await engine.health()
        assert h["status"] == "unavailable"

    @pytest.mark.asyncio
    async def test_load_without_funasr(self) -> None:
        """未安装 funasr 时加载应优雅失败。"""
        engine = SenseVoiceASREngine()
        # 模拟 funasr 不可用
        with patch.dict("sys.modules", {"funasr": None}):
            await engine.load()
        assert engine.model_loaded is False

    @pytest.mark.asyncio
    async def test_transcribe_without_model(self) -> None:
        """模型未加载时 transcribe 返回空结果。"""
        engine = SenseVoiceASREngine()
        result = await engine.transcribe(b"fake-audio", format="wav")
        assert result["text"] == ""
        assert result["model"] == engine._model_name

    @pytest.mark.asyncio
    async def test_normalize_language(self) -> None:
        """语言归一化应正确映射。"""
        engine = SenseVoiceASREngine()
        assert engine._normalize_language(None) == "auto"
        assert engine._normalize_language("") == "auto"
        assert engine._normalize_language("auto") == "auto"
        assert engine._normalize_language("zh") == "zh"
        assert engine._normalize_language("en") == "en"
        assert engine._normalize_language("ja") == "ja"
        assert engine._normalize_language("ko") == "ko"
        assert engine._normalize_language("yue") == "yue"
        # 不支持的应回退 auto
        assert engine._normalize_language("zz") == "auto"

    @pytest.mark.asyncio
    async def test_unload(self) -> None:
        """卸载应清理模型引用。"""
        engine = SenseVoiceASREngine()
        await engine.unload()
        assert engine.model_loaded is False
        assert engine._model is None

    @pytest.mark.asyncio
    async def test_close(self) -> None:
        """close 应调用 unload。"""
        engine = SenseVoiceASREngine()
        await engine.close()
        assert engine.model_loaded is False


# ── 测试端点逻辑 ──────────────────────────────────────────────────────────


class TestEndpoints:
    def test_ensure_loaded_returns_503_when_engine_none(self) -> None:
        """引擎未初始化时返回 503。"""
        with pytest.raises(Exception):
            # 模拟 asr_engine 为 None 的场景
            with patch("src.endpoints.asr.asr_engine", None):
                _ensure_loaded()

    @pytest.mark.asyncio
    async def test_asr_endpoint_no_engine(self) -> None:
        """ASR 引擎未就绪时应有适当的错误响应状态码。"""
        from fastapi import HTTPException

        with patch("src.endpoints.asr.asr_engine", None):
            with pytest.raises(HTTPException) as exc:
                _ensure_loaded()
            assert exc.value.status_code == 503


# ── 集成测试（需要实际模型，标记为跳过） ────────────────────────


@pytest.mark.skip(reason="需要真实 SenseVoiceSmall 模型和 FunASR 环境")
class TestIntegrationASR:
    """集成测试 — 需要实际模型文件。"""

    @pytest.mark.asyncio
    async def test_transcribe_wav(self) -> None:
        """转写 WAV 音频。"""

        engine = SenseVoiceASREngine()
        await engine.load()
        assert engine.model_loaded

        audio = _make_test_wav()
        result = await engine.transcribe(audio, format="wav")
        assert isinstance(result["text"], str)

    @pytest.mark.asyncio
    async def test_transcribe_with_language_hint(self) -> None:
        """指定语言提示的转写。"""
        engine = SenseVoiceASREngine()
        await engine.load()

        audio = _make_test_wav()
        result = await engine.transcribe(audio, format="wav", language="zh")
        assert isinstance(result["text"], str)