test(03-01): add ASRInput and ASROutput schemas
- Add ASRInput schema with video_path and content_type fields - Add ASROutput schema with tool/has_subtitle/output_files/segments/text/duration_seconds/speakers - Create test_processors directory with conftest fixtures - Add comprehensive schema tests (8 test cases) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Tests for kb_pipeline.processors module"""
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Shared fixtures for processor tests"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_whisper_path(tmp_path):
|
||||
"""Mock whisper.cpp path"""
|
||||
whisper_dir = tmp_path / "whisper.cpp"
|
||||
whisper_dir.mkdir()
|
||||
(whisper_dir / "whisper-cli").write_text("")
|
||||
models_dir = whisper_dir / "models"
|
||||
models_dir.mkdir()
|
||||
(models_dir / "ggml-base.bin").write_text("")
|
||||
return str(whisper_dir)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_bilibili_api():
|
||||
"""Mock B站 API 响应"""
|
||||
with patch("aiohttp.ClientSession.get") as mock_get:
|
||||
yield mock_get
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_video(tmp_path):
|
||||
"""创建测试视频文件"""
|
||||
video_path = tmp_path / "test_video.mp4"
|
||||
video_path.write_text("fake video content")
|
||||
return str(video_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_subtitle_response():
|
||||
"""B站字幕 API 响应示例"""
|
||||
return {
|
||||
"data": {
|
||||
"subtitle": {
|
||||
"list": [{
|
||||
"lan": "ai-zh",
|
||||
"lan_doc": "中文（自动生成）",
|
||||
"subtitle_url": "//example.com/subtitle.json"
|
||||
}]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_subtitle_content():
|
||||
"""字幕内容示例"""
|
||||
return {
|
||||
"body": [
|
||||
{"from": 0.0, "to": 3.2, "content": "这是第一段字幕"},
|
||||
{"from": 3.5, "to": 6.8, "content": "这是第二段字幕"},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_whisper_output(tmp_path):
|
||||
"""Whisper 输出示例"""
|
||||
output_base = tmp_path / "test_video"
|
||||
return {
|
||||
"txt": str(output_base.with_suffix(".txt")),
|
||||
"srt": str(output_base.with_suffix(".srt")),
|
||||
"json": str(output_base.with_suffix(".json")),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_state_manager():
|
||||
"""Mock WorkflowStateManager"""
|
||||
manager = MagicMock()
|
||||
manager.load.return_value = None
|
||||
manager.save = MagicMock()
|
||||
return manager
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_settings():
|
||||
"""Mock KBSettings without HF_TOKEN"""
|
||||
settings = MagicMock()
|
||||
settings.hf_token.get_secret_value.return_value = ""
|
||||
return settings
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_settings_with_hf_token():
|
||||
"""Mock KBSettings with HF_TOKEN"""
|
||||
settings = MagicMock()
|
||||
settings.hf_token.get_secret_value.return_value = "test_hf_token_123"
|
||||
return settings
|
||||
@@ -0,0 +1,617 @@
|
||||
"""Tests for ASRProcessor
|
||||
|
||||
测试 ASR 处理器的 Schema 定义、官方字幕下载、Whisper 集成和状态记录功能。
|
||||
"""
|
||||
|
||||
import os
|
||||
import pytest
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch, mock_open
|
||||
|
||||
from pydantic import ValidationError
|
||||
|
||||
# 测试将在实现后导入
|
||||
# from kb_pipeline.schemas import ASRInput, ASROutput
|
||||
# from kb_pipeline.processors.asr import ASRProcessor
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Schema Tests (Task 1)
|
||||
# =============================================================================
|
||||
|
||||
class TestASRInputSchema:
|
||||
"""ASRInput Schema 测试"""
|
||||
|
||||
def test_asr_input_inherits_step_input(self):
|
||||
"""ASRInput 应继承 StepInput 的字段"""
|
||||
from kb_pipeline.schemas import ASRInput, StepInput
|
||||
|
||||
# 创建实例
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path="/path/to/video.mp4"
|
||||
)
|
||||
|
||||
# 验证继承
|
||||
assert isinstance(input_data, StepInput)
|
||||
assert input_data.content_id == "BV1xx"
|
||||
assert input_data.platform == "bilibili"
|
||||
|
||||
def test_asr_input_has_video_path(self):
|
||||
"""ASRInput 必须包含 video_path 字段"""
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path="/path/to/video.mp4"
|
||||
)
|
||||
|
||||
assert hasattr(input_data, "video_path")
|
||||
assert input_data.video_path == "/path/to/video.mp4"
|
||||
|
||||
def test_asr_input_content_type_default(self):
|
||||
"""ASRInput content_type 默认为 'video'"""
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path="/path/to/video.mp4"
|
||||
)
|
||||
|
||||
assert input_data.content_type == "video"
|
||||
|
||||
|
||||
class TestASROutputSchema:
|
||||
"""ASROutput Schema 测试"""
|
||||
|
||||
def test_asr_output_inherits_step_output(self):
|
||||
"""ASROutput 应继承 StepOutput 的字段"""
|
||||
from kb_pipeline.schemas import ASROutput, StepOutput
|
||||
|
||||
output = ASROutput(
|
||||
success=True,
|
||||
tool="subtitle",
|
||||
has_subtitle=True,
|
||||
output_files={"txt": "/path/to/txt", "srt": "/path/to/srt", "json": "/path/to/json"},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=120.5
|
||||
)
|
||||
|
||||
assert isinstance(output, StepOutput)
|
||||
assert output.success is True
|
||||
|
||||
def test_asr_output_tool_literal(self):
|
||||
"""ASROutput tool 字段只能是 'subtitle' | 'whisper' | 'whisper-diarization'"""
|
||||
from kb_pipeline.schemas import ASROutput
|
||||
|
||||
# 有效值
|
||||
for tool in ["subtitle", "whisper", "whisper-diarization"]:
|
||||
output = ASROutput(
|
||||
success=True,
|
||||
tool=tool,
|
||||
has_subtitle=True,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0
|
||||
)
|
||||
assert output.tool == tool
|
||||
|
||||
def test_asr_output_tool_invalid(self):
|
||||
"""ASROutput tool 字段无效值应抛出 ValidationError"""
|
||||
from kb_pipeline.schemas import ASROutput
|
||||
|
||||
with pytest.raises(ValidationError):
|
||||
ASROutput(
|
||||
success=True,
|
||||
tool="invalid_tool",
|
||||
has_subtitle=True,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0
|
||||
)
|
||||
|
||||
def test_asr_output_has_required_fields(self):
|
||||
"""ASROutput 必须包含所有必需字段"""
|
||||
from kb_pipeline.schemas import ASROutput
|
||||
|
||||
output = ASROutput(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={"txt": "a.txt", "srt": "a.srt", "json": "a.json"},
|
||||
segments=[{"start": 0.0, "end": 3.0, "text": "test"}],
|
||||
text="test content",
|
||||
duration_seconds=10.5
|
||||
)
|
||||
|
||||
assert output.tool == "whisper"
|
||||
assert output.has_subtitle is False
|
||||
assert "txt" in output.output_files
|
||||
assert len(output.segments) == 1
|
||||
assert output.text == "test content"
|
||||
assert output.duration_seconds == 10.5
|
||||
|
||||
def test_asr_output_speakers_optional(self):
|
||||
"""ASROutput speakers 字段是可选的"""
|
||||
from kb_pipeline.schemas import ASROutput
|
||||
|
||||
# 无 speakers
|
||||
output1 = ASROutput(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0
|
||||
)
|
||||
assert output1.speakers is None
|
||||
|
||||
# 有 speakers
|
||||
output2 = ASROutput(
|
||||
success=True,
|
||||
tool="whisper-diarization",
|
||||
has_subtitle=False,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0,
|
||||
speakers=["SPEAKER_00", "SPEAKER_01"]
|
||||
)
|
||||
assert output2.speakers == ["SPEAKER_00", "SPEAKER_01"]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Processor Tests (Task 2-5)
|
||||
# =============================================================================
|
||||
|
||||
class TestASRProcessorInit:
|
||||
"""ASRProcessor 初始化测试"""
|
||||
|
||||
def test_asr_processor_init_success(self, mock_settings):
|
||||
"""ASRProcessor 初始化成功"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
assert processor is not None
|
||||
assert processor.STEP_NAME == "asr"
|
||||
|
||||
def test_asr_processor_with_state_manager(self, mock_settings, mock_state_manager):
|
||||
"""ASRProcessor 可接受 state_manager 参数"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor(state_manager=mock_state_manager)
|
||||
assert processor.state_manager is mock_state_manager
|
||||
|
||||
|
||||
class TestOfficialSubtitle:
|
||||
"""官方字幕下载测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_official_subtitle_found(self, mock_settings, sample_subtitle_response, sample_subtitle_content, tmp_path):
|
||||
"""_try_official_subtitle 找到字幕时返回 ASROutput"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
# Mock API 调用
|
||||
with patch("aiohttp.ClientSession.get") as mock_get:
|
||||
mock_resp = AsyncMock()
|
||||
mock_resp.json = AsyncMock(side_effect=[
|
||||
{"data": [{"cid": 123456}]}, # pagelist 响应
|
||||
sample_subtitle_response, # subtitle 响应
|
||||
sample_subtitle_content, # 字幕内容
|
||||
])
|
||||
mock_resp.__aenter__.return_value = mock_resp
|
||||
mock_get.return_value = mock_resp
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
result = await processor._try_official_subtitle(input_data)
|
||||
|
||||
assert result is not None
|
||||
assert result.tool == "subtitle"
|
||||
assert result.has_subtitle is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_official_subtitle_not_found(self, mock_settings, tmp_path):
|
||||
"""_try_official_subtitle 无字幕时返回 None"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
# Mock API 调用 - 无字幕
|
||||
with patch("aiohttp.ClientSession.get") as mock_get:
|
||||
mock_resp = AsyncMock()
|
||||
mock_resp.json = AsyncMock(side_effect=[
|
||||
{"data": [{"cid": 123456}]}, # pagelist 响应
|
||||
{"data": {"subtitle": {"list": []}}}, # 无字幕
|
||||
])
|
||||
mock_resp.__aenter__.return_value = mock_resp
|
||||
mock_get.return_value = mock_resp
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
result = await processor._try_official_subtitle(input_data)
|
||||
|
||||
# 无字幕返回 None（不是抛异常）
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_tries_subtitle_first(self, mock_settings, tmp_path):
|
||||
"""process() 方法优先尝试官方字幕"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
# Mock _try_official_subtitle 返回成功
|
||||
with patch.object(processor, "_try_official_subtitle") as mock_subtitle:
|
||||
mock_subtitle.return_value = MagicMock(tool="subtitle", success=True)
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
await processor.process(input_data)
|
||||
|
||||
# 应该调用 _try_official_subtitle
|
||||
mock_subtitle.assert_called_once()
|
||||
|
||||
|
||||
class TestWhisperBasic:
|
||||
"""Whisper.cpp 基础测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_whisper_basic(self, mock_settings, mock_whisper_path, tmp_path):
|
||||
"""_run_whisper 调用成功返回 ASROutput"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
# 创建 Whisper 输出文件
|
||||
txt_path = tmp_path / "test.txt"
|
||||
srt_path = tmp_path / "test.srt"
|
||||
json_path = tmp_path / "test.json"
|
||||
txt_path.write_text("transcription text")
|
||||
srt_path.write_text("1\n00:00:00,000 --> 00:00:03,000\ntest\n")
|
||||
json_path.write_text('{"transcription": [{"text": "test"}]}')
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_whisper_cpp_path", return_value=mock_whisper_path):
|
||||
with patch("subprocess.run") as mock_run:
|
||||
mock_run.return_value = MagicMock(returncode=0)
|
||||
|
||||
with patch("kb_pipeline.processors.asr.convert_to_wav") as mock_convert:
|
||||
mock_convert.return_value = str(tmp_path / "test_audio.wav")
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
result = await processor._run_whisper_basic(input_data)
|
||||
|
||||
assert result.success is True
|
||||
assert result.tool == "whisper"
|
||||
assert result.has_subtitle is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_temp_file_cleanup(self, mock_settings, mock_whisper_path, tmp_path):
|
||||
"""临时音频文件处理后自动清理"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
audio_path = tmp_path / "test_audio.wav"
|
||||
audio_path.write_text("fake audio")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_whisper_cpp_path", return_value=mock_whisper_path):
|
||||
with patch("subprocess.run") as mock_run:
|
||||
mock_run.return_value = MagicMock(returncode=0)
|
||||
|
||||
with patch("kb_pipeline.processors.asr.convert_to_wav", return_value=str(audio_path)):
|
||||
with patch.object(processor, "_parse_whisper_output") as mock_parse:
|
||||
mock_parse.return_value = MagicMock(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0
|
||||
)
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
await processor._run_whisper_basic(input_data)
|
||||
|
||||
# 临时文件应该被删除
|
||||
assert not audio_path.exists()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_whisper_path_not_found(self, mock_settings, tmp_path):
|
||||
"""Whisper 路径不存在时抛出 RuntimeError"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor, WhisperNotFoundError
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_whisper_cpp_path", side_effect=WhisperNotFoundError("Not found")):
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await processor._run_whisper_basic(input_data)
|
||||
|
||||
|
||||
class TestDiarization:
|
||||
"""说话者分离测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_hf_token_uses_basic_whisper(self, mock_settings, mock_whisper_path, tmp_path):
|
||||
"""HF_TOKEN 未配置时使用普通 Whisper"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor()
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_whisper_cpp_path", return_value=mock_whisper_path):
|
||||
with patch("subprocess.run") as mock_run:
|
||||
mock_run.return_value = MagicMock(returncode=0)
|
||||
|
||||
with patch("kb_pipeline.processors.asr.convert_to_wav") as mock_convert:
|
||||
mock_convert.return_value = str(tmp_path / "test_audio.wav")
|
||||
|
||||
with patch.object(processor, "_parse_whisper_output") as mock_parse:
|
||||
mock_parse.return_value = MagicMock(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0
|
||||
)
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
result = await processor._run_whisper(input_data)
|
||||
|
||||
# 应该是普通 whisper，不是 diarization
|
||||
assert result.tool == "whisper"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_diarization_with_hf_token(self, mock_settings_with_hf_token, mock_whisper_path, tmp_path):
|
||||
"""HF_TOKEN 配置时自动启用说话者分离"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings_with_hf_token):
|
||||
processor = ASRProcessor()
|
||||
|
||||
with patch.object(processor, "_run_whisper_with_diarization") as mock_diarization:
|
||||
mock_diarization.return_value = MagicMock(
|
||||
success=True,
|
||||
tool="whisper-diarization",
|
||||
has_subtitle=False,
|
||||
output_files={},
|
||||
segments=[],
|
||||
text="",
|
||||
duration_seconds=0.0,
|
||||
speakers=["SPEAKER_00"]
|
||||
)
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
result = await processor._run_whisper(input_data)
|
||||
|
||||
# 应该调用 diarization
|
||||
mock_diarization.assert_called_once()
|
||||
assert result.tool == "whisper-diarization"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_diarization_result_has_speaker(self, mock_settings_with_hf_token, mock_whisper_path, tmp_path):
|
||||
"""说话者分离结果包含 speaker 字段"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings_with_hf_token):
|
||||
processor = ASRProcessor()
|
||||
|
||||
with patch.object(processor, "_run_whisper_basic") as mock_basic:
|
||||
mock_basic.return_value = MagicMock(
|
||||
success=True,
|
||||
tool="whisper",
|
||||
has_subtitle=False,
|
||||
output_files={},
|
||||
segments=[{"start": 0.0, "end": 3.0, "text": "hello"}],
|
||||
text="hello",
|
||||
duration_seconds=3.0
|
||||
)
|
||||
|
||||
with patch("kb_pipeline.processors.asr.run_diarization") as mock_dia:
|
||||
mock_dia.return_value = {
|
||||
"SPEAKER_00": [{"start": 0.0, "end": 3.0}]
|
||||
}
|
||||
|
||||
with patch("kb_pipeline.processors.asr.assign_speakers_to_segments") as mock_assign:
|
||||
mock_assign.return_value = [
|
||||
{"start": 0.0, "end": 3.0, "text": "hello", "speaker": "SPEAKER_00"}
|
||||
]
|
||||
|
||||
result = await processor._run_whisper_with_diarization(
|
||||
MagicMock(video_path=str(video_path), content_id="BV1xx", platform="bilibili")
|
||||
)
|
||||
|
||||
# 结果应包含 speakers 列表
|
||||
assert result.speakers is not None
|
||||
assert len(result.speakers) > 0
|
||||
|
||||
|
||||
class TestStateRecording:
|
||||
"""状态记录测试"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_recording_success(self, mock_settings, mock_state_manager, tmp_path):
|
||||
"""ASR 成功时状态记录为 completed"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput, WorkflowState
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
# 创建 mock state
|
||||
mock_state = MagicMock(spec=WorkflowState)
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor(state_manager=mock_state_manager)
|
||||
mock_state_manager.load.return_value = mock_state
|
||||
|
||||
with patch.object(processor, "_try_official_subtitle") as mock_sub:
|
||||
mock_sub.return_value = MagicMock(
|
||||
success=True,
|
||||
tool="subtitle",
|
||||
segments=[1, 2, 3]
|
||||
)
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
await processor.process(input_data)
|
||||
|
||||
# 应该调用 mark_step_completed
|
||||
mock_state.mark_step_completed.assert_called_once()
|
||||
call_args = mock_state.mark_step_completed.call_args
|
||||
assert call_args[0][0] == "asr"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_recording_failure(self, mock_settings, mock_state_manager, tmp_path):
|
||||
"""ASR 失败时状态记录为 failed"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput, WorkflowState
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
mock_state = MagicMock(spec=WorkflowState)
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor(state_manager=mock_state_manager)
|
||||
mock_state_manager.load.return_value = mock_state
|
||||
|
||||
with patch.object(processor, "_try_official_subtitle", side_effect=Exception("Network error")):
|
||||
with patch.object(processor, "_run_whisper", side_effect=Exception("Whisper error")):
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await processor.process(input_data)
|
||||
|
||||
# 应该调用 mark_step_failed
|
||||
mock_state.mark_step_failed.assert_called_once()
|
||||
call_args = mock_state.mark_step_failed.call_args
|
||||
assert call_args[0][0] == "asr"
|
||||
assert "Whisper error" in call_args[0][1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_state_no_manager(self, mock_settings, tmp_path):
|
||||
"""无 state_manager 时不记录状态"""
|
||||
from kb_pipeline.processors.asr import ASRProcessor
|
||||
from kb_pipeline.schemas import ASRInput
|
||||
|
||||
video_path = tmp_path / "test.mp4"
|
||||
video_path.write_text("fake video")
|
||||
|
||||
with patch("kb_pipeline.processors.asr.get_settings", return_value=mock_settings):
|
||||
processor = ASRProcessor(state_manager=None)
|
||||
|
||||
with patch.object(processor, "_try_official_subtitle") as mock_sub:
|
||||
mock_sub.return_value = MagicMock(success=True, tool="subtitle", segments=[])
|
||||
|
||||
input_data = ASRInput(
|
||||
content_id="BV1xx",
|
||||
platform="bilibili",
|
||||
video_path=str(video_path)
|
||||
)
|
||||
|
||||
# 应该不抛异常
|
||||
await processor.process(input_data)
|
||||
Reference in New Issue
Block a user