“ydy0615”
264 lines
8.5 KiB
Python
264 lines
8.5 KiB
Python
import os
|
|
import sys
|
|
import base64
|
|
import io
|
|
import types
|
|
import wave
|
|
import pytest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
import numpy as np
|
|
|
|
BACKEND_DIR = Path(__file__).resolve().parents[1]
|
|
if str(BACKEND_DIR) not in sys.path:
|
|
sys.path.insert(0, str(BACKEND_DIR))
|
|
|
|
|
|
def _make_mlx_stub():
|
|
"""Create minimal MLX stub for testing without Apple Silicon"""
|
|
mlx = types.SimpleNamespace()
|
|
mlx.core = types.SimpleNamespace()
|
|
mx_array = type('mx.array', (), {'item': lambda self: 1})
|
|
mlx.core.array = mx_array
|
|
|
|
def mock_load(path):
|
|
return MagicMock()
|
|
mlx.core.load = mock_load # type: ignore
|
|
|
|
mlx.nn = types.SimpleNamespace()
|
|
return mlx
|
|
|
|
|
|
def _make_mlx_audio_stub():
|
|
"""Create minimal mlx-audio stub"""
|
|
stt = types.SimpleNamespace()
|
|
stt.utils = types.SimpleNamespace()
|
|
|
|
def mock_load(path): # type: ignore
|
|
model = MagicMock()
|
|
output = types.SimpleNamespace()
|
|
output.text = "识别结果"
|
|
output.language = "zh-CN"
|
|
model.generate = MagicMock(return_value=output)
|
|
return model
|
|
|
|
stt.utils.load = mock_load # type: ignore
|
|
|
|
qwen3_asr_mod = types.SimpleNamespace()
|
|
qwen3_asr_mod.Qwen3ASRModel = type('Qwen3ASRModel', (), {})
|
|
qwen3_asr_mod.ForcedAlignerModel = type('ForcedAlignerModel', (), {})
|
|
stt.models = types.SimpleNamespace() # type: ignore
|
|
stt.models.qwen3_asr = qwen3_asr_mod # type: ignore
|
|
|
|
audio = types.SimpleNamespace()
|
|
audio.stt = stt # type: ignore
|
|
return audio
|
|
|
|
|
|
def _reload_tts_asr_with_mocks():
|
|
"""Reload tts_asr with mocked MLX dependencies"""
|
|
for mod_name in list(sys.modules.keys()):
|
|
if 'tts_asr' in mod_name or 'mlx' in mod_name:
|
|
del sys.modules[mod_name]
|
|
|
|
mlx_stub = _make_mlx_stub()
|
|
sys.modules['mlx'] = mlx_stub # type: ignore
|
|
sys.modules['mlx.core'] = mlx_stub.core # type: ignore
|
|
sys.modules['mlx.nn'] = mlx_stub.nn # type: ignore
|
|
|
|
audio_stub = _make_mlx_audio_stub()
|
|
sys.modules['mlx-audio'] = audio_stub # type: ignore
|
|
sys.modules['mlx_audio'] = audio_stub # type: ignore
|
|
sys.modules['mlx_audio.stt'] = audio_stub.stt # type: ignore
|
|
sys.modules['mlx_audio.stt.utils'] = audio_stub.stt.utils # type: ignore
|
|
sys.modules['mlx_audio.stt.models'] = audio_stub.stt.models # type: ignore
|
|
sys.modules['mlx_audio.stt.models.qwen3_asr'] = audio_stub.stt.models.qwen3_asr # type: ignore
|
|
|
|
import tts_asr
|
|
return tts_asr, audio_stub
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _clean_env():
|
|
"""Clean ASR-related env vars before/after each test"""
|
|
saved = {}
|
|
for k in ['HF_ENDPOINT']:
|
|
saved[k] = os.environ.get(k)
|
|
if k in os.environ:
|
|
del os.environ[k]
|
|
yield
|
|
for k, v in saved.items():
|
|
if v is not None:
|
|
os.environ[k] = v # type: ignore
|
|
|
|
|
|
def _make_wav_bytes(sr=16000, duration_sec=1.0, channels=1):
|
|
"""Helper: generate WAV bytes as base64"""
|
|
samples = int(sr * duration_sec)
|
|
audio = np.random.randint(-32768, 32767, size=samples * channels, dtype=np.int16)
|
|
buf = io.BytesIO()
|
|
with wave.open(buf, 'wb') as wf:
|
|
wf.setnchannels(channels)
|
|
wf.setsampwidth(2)
|
|
wf.setframerate(sr)
|
|
wf.writeframes(audio.tobytes())
|
|
return base64.b64encode(buf.getvalue()).decode()
|
|
|
|
|
|
class TestASRLazyLoading:
|
|
"""测试 ASR 模型懒加载"""
|
|
|
|
def test_ensure_asr_loads_on_call(self):
|
|
tts, audio_stub = _reload_tts_asr_with_mocks()
|
|
assert tts._asr_model is None
|
|
|
|
model = tts._ensure_asr_model()
|
|
assert model is not None
|
|
|
|
def test_ensure_align_loads_on_call(self):
|
|
tts, audio_stub = _reload_tts_asr_with_mocks()
|
|
assert tts._align_model is None
|
|
|
|
model = tts._ensure_align_model()
|
|
assert model is not None
|
|
|
|
|
|
class TestASREndpoint:
|
|
"""测试 ASR 端点逻辑"""
|
|
|
|
def test_asr_basic_recognition(self, fastapi_testclient=None):
|
|
"""ASR 端点应正确返回识别结果"""
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
|
|
# Mock the model to return known values
|
|
tts._asr_model = MagicMock()
|
|
output = types.SimpleNamespace()
|
|
output.text = "你好世界"
|
|
output.language = "zh-CN"
|
|
tts._asr_model.generate.return_value = output
|
|
|
|
wav_b64 = _make_wav_bytes()
|
|
req = tts.ASRRequest(audio_base64=wav_b64)
|
|
|
|
# Call generate directly (simulating endpoint logic)
|
|
audio_bytes = base64.b64decode(req.audio_base64)
|
|
wav_buffer = io.BytesIO(audio_bytes)
|
|
with wave.open(wav_buffer, 'rb') as wf:
|
|
raw = wf.readframes(wf.getnframes())
|
|
arr = np.frombuffer(raw, dtype=np.int16)
|
|
arr = arr.astype(np.float32) / 32768.0
|
|
|
|
result = tts._asr_model.generate(arr, language=req.language)
|
|
assert result.text == "你好世界"
|
|
|
|
def test_asr_stereo_to_mono(self):
|
|
"""立体声音频应被正确转换为单声道"""
|
|
wav_b64 = _make_wav_bytes(channels=2)
|
|
|
|
audio_bytes = base64.b64decode(wav_b64)
|
|
wav_buffer = io.BytesIO(audio_bytes)
|
|
with wave.open(wav_buffer, 'rb') as wf:
|
|
assert wf.getnchannels() == 2
|
|
n_frames = wf.getnframes()
|
|
raw_data = wf.readframes(n_frames)
|
|
audio_array = np.frombuffer(raw_data, dtype=np.int16)
|
|
|
|
# Convert to mono
|
|
audio_array = np.mean(audio_array.reshape(-1, 2), axis=1)
|
|
assert audio_array.ndim == 1
|
|
|
|
def test_asr_resample_to_16k(self):
|
|
"""非 16kHz 音频应被重采样"""
|
|
wav_b64 = _make_wav_bytes(sr=48000, duration_sec=0.5)
|
|
|
|
audio_bytes = base64.b64decode(wav_b64)
|
|
wav_buffer = io.BytesIO(audio_bytes)
|
|
with wave.open(wav_buffer, 'rb') as wf:
|
|
assert wf.getframerate() == 48000
|
|
|
|
def test_asr_44100_resample(self):
|
|
"""44.1kHz 常见采样率应被重采样到 16k"""
|
|
wav_b64 = _make_wav_bytes(sr=44100, duration_sec=1.0)
|
|
|
|
audio_bytes = base64.b64decode(wav_b64)
|
|
wav_buffer = io.BytesIO(audio_bytes)
|
|
with wave.open(wav_buffer, 'rb') as wf:
|
|
framerate = wf.getframerate()
|
|
n_frames = wf.getnframes()
|
|
raw_data = wf.readframes(n_frames)
|
|
audio_array = np.frombuffer(raw_data, dtype=np.int16)
|
|
|
|
# Simulate resample calculation
|
|
if framerate != 16000:
|
|
n_samples = int(len(audio_array) * 16000 / framerate)
|
|
else:
|
|
n_samples = len(audio_array)
|
|
|
|
expected_16k_samples = int(1.0 * 16000)
|
|
assert abs(n_samples - expected_16k_samples) < 2
|
|
|
|
|
|
class TestASRModelDownload:
|
|
"""测试 ASR 模型下载路径"""
|
|
|
|
def test_load_asr_from_path_success(self):
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
|
|
with patch('backend.tts_asr.snapshot_download', return_value='/fake/asr'): # type: ignore
|
|
tts._load_asr_models()
|
|
|
|
assert tts._asr_model is not None
|
|
|
|
def test_load_asr_skips_without_mlx(self):
|
|
"""不注入 MLX stub 时应跳过 ASR"""
|
|
for mod_name in list(sys.modules.keys()):
|
|
if 'tts_asr' in mod_name or 'mlx' in mod_name:
|
|
del sys.modules[mod_name]
|
|
|
|
import tts_asr # noqa: F811
|
|
assert tts_asr.Qwen3ASRModel is None
|
|
|
|
def test_load_align_from_path(self):
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
|
|
with patch('backend.tts_asr.snapshot_download', return_value='/fake/align'): # type: ignore
|
|
tts._load_asr_models()
|
|
|
|
assert tts._align_model is not None
|
|
|
|
|
|
class TestModelConstants:
|
|
"""测试模型 ID 常量"""
|
|
|
|
def test_asr_model_id(self):
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
assert "aufklarer" in tts.ASR_MODEL_ID_MS
|
|
|
|
def test_align_model_id(self):
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
assert "ForcedAligner" in tts.ALIGN_MODEL_ID_MS
|
|
|
|
def test_tts_model_id(self):
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
assert "Qwen3-TTS" in tts.MODEL_ID_MS
|
|
|
|
|
|
class TestHFEndpointMirror:
|
|
"""测试镜像站配置"""
|
|
|
|
def test_hf_endpoint_set(self):
|
|
tts, _ = _reload_tts_asr_with_mocks()
|
|
assert os.environ.get("HF_ENDPOINT") == "https://hf-mirror.com"
|
|
|
|
def test_hf_endpoint_default(self):
|
|
"""即使环境变量未设置，模块也应默认设置镜像"""
|
|
for mod_name in list(sys.modules.keys()):
|
|
if 'tts_asr' in mod_name or 'mlx' in mod_name:
|
|
del sys.modules[mod_name]
|
|
|
|
if "HF_ENDPOINT" in os.environ:
|
|
del os.environ["HF_ENDPOINT"]
|
|
|
|
import tts_asr # noqa: F811
|
|
assert os.environ.get("HF_ENDPOINT") == "https://hf-mirror.com"
|