diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ccc203a --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,62 @@ +"""共享测试 fixtures 和配置""" + +from typing import Optional +from unittest.mock import Mock + +import pytest + + +class MockAIClient: + """模拟 AI 客户端""" + + def __init__(self, response: str = "Mock AI response"): + self.response = response + + def generate( + self, + prompt: str, + max_tokens: int = 16000, + disable_reasoning: Optional[bool] = None + ) -> str: + """生成模拟响应""" + return self.response + + +@pytest.fixture +def mock_ai_client(): + """提供模拟 AI 客户端""" + return MockAIClient() + + +@pytest.fixture +def sample_segments(): + """提供示例字幕分段数据""" + return [ + {"index": 1, "start": "00:00:00,000", "end": "00:00:05,000", "text": "大家好,欢迎来到我的频道"}, + {"index": 2, "start": "00:00:05,000", "end": "00:00:10,000", "text": "今天我们要讲的是机器学习的基础知识"}, + {"index": 3, "start": "00:00:10,000", "end": "00:00:15,000", "text": "首先,什么是机器学习"}, + {"index": 4, "start": "00:00:15,000", "end": "00:00:20,000", "text": "机器学习是人工智能的一个分支"}, + {"index": 5, "start": "00:00:20,000", "end": "00:00:25,000", "text": "它让计算机能够从数据中学习"}, + ] + + +@pytest.fixture +def sample_metadata(): + """提供示例视频元数据""" + return { + "title": "机器学习入门教程", + "category": "教育", + "tags": ["编程", "AI", "教程"], + "description": "从零开始学习机器学习" + } + + +@pytest.fixture +def sample_metadata_movie(): + """提供电影类型示例元数据""" + return { + "title": "星际穿越", + "category": "影视", + "tags": ["科幻", "电影"], + "description": "一部关于太空探索的科幻电影" + } diff --git a/tests/test_semantic_engine.py b/tests/test_semantic_engine.py new file mode 100644 index 0000000..928dcf3 --- /dev/null +++ b/tests/test_semantic_engine.py @@ -0,0 +1,129 @@ +"""测试语义分段引擎""" + +import pytest + +from kb_workflow.srt.semantic_engine import ( + VideoSegment, + SegmentationResult, + SemanticSegmentationEngine +) + + +class TestVideoTypeDetection: + """测试视频类型识别功能""" + + def test_video_type_detection_from_title(self): + """测试从标题推断视频类型""" + # 知识类 - 从标题推断 + metadata_edu = {"title": "Python编程教程", "category": "教育"} + engine = SemanticSegmentationEngine() + video_type = engine._infer_video_type(metadata_edu) + assert video_type == "知识类" + + # 知识类 - 包含"课程" + metadata_course = {"title": "机器学习课程", "category": ""} + video_type = engine._infer_video_type(metadata_course) + assert video_type == "知识类" + + def test_video_type_detection_from_category(self): + """测试从分类推断视频类型""" + # 故事类 - 从分类推断 + metadata_movie = {"title": "星际穿越", "category": "影视"} + engine = SemanticSegmentationEngine() + video_type = engine._infer_video_type(metadata_movie) + assert video_type == "故事类" + + # 访谈类 - 从标签推断 + metadata_interview = {"title": "名人访谈", "tags": ["访谈", "对话"]} + video_type = engine._infer_video_type(metadata_interview) + assert video_type == "访谈类" + + def test_video_type_fallback_to_default(self): + """测试元数据缺失时使用默认类型""" + # 空元数据 + engine = SemanticSegmentationEngine() + video_type = engine._infer_video_type({}) + assert video_type == "知识类" # 默认类型 + + # 无效元数据 + video_type = engine._infer_video_type({"title": "", "category": ""}) + assert video_type == "知识类" + + +class TestVideoTypeInPrompt: + """测试视频类型在 Prompt 中的使用""" + + def test_video_type_in_prompt(self): + """测试 Prompt 包含视频类型信息""" + engine = SemanticSegmentationEngine() + metadata = {"title": "机器学习教程", "category": "教育"} + + prompt = engine._build_prompt("测试内容", metadata) + + assert "视频类型" in prompt or "知识类" in prompt + assert "机器学习教程" in prompt + + def test_prompt_format(self): + """测试 Prompt 格式正确""" + engine = SemanticSegmentationEngine() + content = "这是测试内容" + metadata = {"title": "测试视频"} + + prompt = engine._build_prompt(content, metadata) + + assert "请分析以下视频字幕" in prompt or "字幕内容" in prompt + assert content in prompt + + +class TestPydanticModels: + """测试 Pydantic 模型验证""" + + def test_video_segment_model(self): + """测试 VideoSegment 模型""" + # 创建完整分段 + segment = VideoSegment( + title="机器学习概述", + content="机器学习是人工智能的一个分支", + start_index=1, + end_index=5 + ) + assert segment.title == "机器学习概述" + assert segment.content == "机器学习是人工智能的一个分支" + assert segment.start_index == 1 + assert segment.end_index == 5 + + # 创建无标题分段(title 可选) + segment_no_title = VideoSegment( + content="测试内容", + start_index=1, + end_index=3 + ) + assert segment_no_title.title is None + + def test_segmentation_result_model(self): + """测试 SegmentationResult 模型""" + segments = [ + VideoSegment(title="第一章", content="内容1", start_index=1, end_index=5), + VideoSegment(title="第二章", content="内容2", start_index=6, end_index=10) + ] + + result = SegmentationResult( + video_type="知识类", + segments=segments, + actual_segment_count=2 + ) + + assert result.video_type == "知识类" + assert len(result.segments) == 2 + assert result.actual_segment_count == 2 + + def test_model_validation(self): + """测试模型字段验证""" + # 缺少必填字段应该抛出异常 + with pytest.raises(Exception): # Pydantic ValidationError + VideoSegment( + title="测试", + # 缺少 content + start_index=1, + end_index=3 + )