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:
+49
-1
@@ -6,7 +6,7 @@ StepStatus、StepRecord、WorkflowState 和 StepDefinition。
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
from pydantic import BaseModel, Field
|
||||||
from typing import Optional, Dict, Literal, Any
|
from typing import Optional, Dict, Literal, Any, List
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
|
|
||||||
@@ -185,3 +185,51 @@ class StepDefinition(BaseModel):
|
|||||||
input_schema: type[StepInput] = Field(..., description="输入 Schema 类型")
|
input_schema: type[StepInput] = Field(..., description="输入 Schema 类型")
|
||||||
output_schema: type[StepOutput] = Field(..., description="输出 Schema 类型")
|
output_schema: type[StepOutput] = Field(..., description="输出 Schema 类型")
|
||||||
dependencies: list[str] = Field(default_factory=list, description="依赖的步骤列表")
|
dependencies: list[str] = Field(default_factory=list, description="依赖的步骤列表")
|
||||||
|
|
||||||
|
|
||||||
|
# =============================================================================
|
||||||
|
# ASR Schemas (Phase 3)
|
||||||
|
# =============================================================================
|
||||||
|
|
||||||
|
class ASRInput(StepInput):
|
||||||
|
"""ASR 步骤输入
|
||||||
|
|
||||||
|
继承 StepInput 的基础字段,添加视频路径和内容类型。
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
video_path: 视频文件路径
|
||||||
|
content_type: 内容类型,默认为 'video'
|
||||||
|
"""
|
||||||
|
video_path: str = Field(..., description="视频文件路径")
|
||||||
|
content_type: str = Field(default="video", description="内容类型")
|
||||||
|
|
||||||
|
|
||||||
|
class ASROutput(StepOutput):
|
||||||
|
"""ASR 步骤输出
|
||||||
|
|
||||||
|
继承 StepOutput 的基础字段,添加 ASR 特定字段。
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
tool: 使用的工具 (subtitle, whisper, whisper-diarization)
|
||||||
|
has_subtitle: 是否有官方字幕
|
||||||
|
output_files: 输出文件路径字典 {txt, srt, json}
|
||||||
|
segments: 带时间戳的分段列表
|
||||||
|
text: 完整文本
|
||||||
|
duration_seconds: 音频时长
|
||||||
|
speakers: 说话者列表(仅说话者分离时有值)
|
||||||
|
"""
|
||||||
|
tool: Literal["subtitle", "whisper", "whisper-diarization"] = Field(
|
||||||
|
..., description="使用的工具"
|
||||||
|
)
|
||||||
|
has_subtitle: bool = Field(..., description="是否有官方字幕")
|
||||||
|
output_files: Dict[str, str] = Field(
|
||||||
|
default_factory=dict, description="输出文件路径 {txt, srt, json}"
|
||||||
|
)
|
||||||
|
segments: List[Dict[str, Any]] = Field(
|
||||||
|
default_factory=list, description="带时间戳的分段列表"
|
||||||
|
)
|
||||||
|
text: str = Field(default="", description="完整文本")
|
||||||
|
duration_seconds: float = Field(default=0.0, description="音频时长")
|
||||||
|
speakers: Optional[List[str]] = Field(
|
||||||
|
None, description="说话者列表(仅说话者分离时有值)"
|
||||||
|
)
|
||||||
|
|||||||
@@ -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