diff --git a/kb_pipeline/schemas.py b/kb_pipeline/schemas.py index 4504a6f..1d3011f 100644 --- a/kb_pipeline/schemas.py +++ b/kb_pipeline/schemas.py @@ -6,7 +6,7 @@ StepStatus、StepRecord、WorkflowState 和 StepDefinition。 """ 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 enum import Enum @@ -185,3 +185,51 @@ class StepDefinition(BaseModel): input_schema: type[StepInput] = Field(..., description="输入 Schema 类型") output_schema: type[StepOutput] = Field(..., description="输出 Schema 类型") 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="说话者列表(仅说话者分离时有值)" + ) diff --git a/tests/test_processors/__init__.py b/tests/test_processors/__init__.py new file mode 100644 index 0000000..3642005 --- /dev/null +++ b/tests/test_processors/__init__.py @@ -0,0 +1 @@ +"""Tests for kb_pipeline.processors module""" diff --git a/tests/test_processors/conftest.py b/tests/test_processors/conftest.py new file mode 100644 index 0000000..688ac6e --- /dev/null +++ b/tests/test_processors/conftest.py @@ -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 diff --git a/tests/test_processors/test_asr.py b/tests/test_processors/test_asr.py new file mode 100644 index 0000000..0f46d0d --- /dev/null +++ b/tests/test_processors/test_asr.py @@ -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)