"""Tests for SRT processing module Tests for: - SRT parser (parse_srt, parse_srt_file) - Semantic segmenter (SemanticSegmenter) - ASR corrector (ASRCorrector) - SRTProcessor (pipeline processor) """ import asyncio import json import pytest from pathlib import Path from dataclasses import dataclass from typing import List, Optional # ============================================================================= # Test Fixtures # ============================================================================= @pytest.fixture def sample_srt_content() -> str: """标准 SRT 格式内容""" return """1 00:00:00,000 --> 00:00:05,000 第一句字幕内容 2 00:00:05,000 --> 00:00:10,000 第二句字幕内容 3 00:00:10,000 --> 00:00:15,000 第三句字幕内容 """ @pytest.fixture def sample_srt_with_special_chars() -> str: """包含特殊字符的 SRT 内容""" return """1 00:00:00,000 --> 00:00:05,000 你好，世界！这是测试。 2 00:00:05,000 --> 00:00:10,000 包含特殊字符：<>&"' 3 00:00:10,000 --> 00:00:15,000 多行 字幕内容 """ @pytest.fixture def invalid_srt_content() -> str: """格式错误的 SRT 内容""" return "这不是有效的 SRT 格式" @pytest.fixture def empty_content() -> str: """空内容""" return "" @pytest.fixture def sample_srt_file(tmp_path: Path, sample_srt_content: str) -> Path: """创建临时 SRT 文件""" srt_file = tmp_path / "test.srt" srt_file.write_text(sample_srt_content, encoding="utf-8") return srt_file # ============================================================================= # Test 1: SRT Parser - parse_srt # ============================================================================= class TestParseSrt: """测试 SRT 解析功能""" def test_parse_standard_srt(self, sample_srt_content: str): """Test 1: 标准的 SRT 格式能被正确解析为 Segment 列表""" from kb_workflow.srt.parser import parse_srt, Segment segments = parse_srt(sample_srt_content) assert len(segments) == 3 assert all(isinstance(s, Segment) for s in segments) # 验证第一个 segment assert segments[0].id == 1 assert segments[0].start_time == 0.0 assert segments[0].end_time == 5.0 assert segments[0].text == "第一句字幕内容" # 验证最后一个 segment assert segments[2].id == 3 assert segments[2].start_time == 10.0 assert segments[2].end_time == 15.0 assert segments[2].text == "第三句字幕内容" def test_parse_srt_with_special_chars(self, sample_srt_with_special_chars: str): """Test 2: 包含特殊字符的字幕能正确处理""" from kb_workflow.srt.parser import parse_srt segments = parse_srt(sample_srt_with_special_chars) assert len(segments) == 3 assert "你好，世界！" in segments[0].text assert "<>&\"'" in segments[1].text # 多行字幕应该合并 assert "多行" in segments[2].text def test_parse_invalid_srt_raises_error(self, invalid_srt_content: str): """Test 3: 空文件或格式错误的文件抛出 ValueError""" from kb_workflow.srt.parser import parse_srt with pytest.raises(ValueError, match="Invalid SRT format"): parse_srt(invalid_srt_content) def test_parse_empty_content_raises_error(self, empty_content: str): """Test 3b: 空内容抛出 ValueError""" from kb_workflow.srt.parser import parse_srt with pytest.raises(ValueError, match="Empty SRT content"): parse_srt(empty_content) # ============================================================================= # Test 1b: SRT Parser - parse_srt_file # ============================================================================= class TestParseSrtFile: """测试从文件解析 SRT""" def test_parse_srt_file(self, sample_srt_file: Path): """从文件解析 SRT""" from kb_workflow.srt.parser import parse_srt_file segments = parse_srt_file(sample_srt_file) assert len(segments) == 3 assert segments[0].text == "第一句字幕内容" def test_parse_nonexistent_file_raises_error(self): """不存在的文件抛出 FileNotFoundError""" from kb_workflow.srt.parser import parse_srt_file with pytest.raises(FileNotFoundError): parse_srt_file(Path("/nonexistent/file.srt")) # ============================================================================= # Test 2: Semantic Segmenter # ============================================================================= @pytest.fixture def short_segments(): """连续短句字幕分段""" from kb_workflow.srt.parser import Segment return [ Segment(id=1, start_time=0.0, end_time=2.0, text="这是第一句短句"), Segment(id=2, start_time=2.0, end_time=4.0, text="这是第二句短句"), Segment(id=3, start_time=4.0, end_time=6.0, text="这是第三句短句"), ] @pytest.fixture def segments_with_gap(): """带有长停顿的字幕分段""" from kb_workflow.srt.parser import Segment return [ Segment(id=1, start_time=0.0, end_time=5.0, text="第一段内容在这里"), Segment(id=2, start_time=5.0, end_time=10.0, text="继续第一段"), # 5秒停顿 Segment(id=3, start_time=15.0, end_time=20.0, text="第二段从这里开始"), Segment(id=4, start_time=20.0, end_time=25.0, text="继续第二段"), ] @pytest.fixture def long_text_segment(): """长文本分段（超过500字符）""" from kb_workflow.srt.parser import Segment long_text = "这是一段很长的文本，" * 50 # 约600字符 return [ Segment(id=1, start_time=0.0, end_time=60.0, text=long_text), ] class TestSemanticSegmenter: """测试语义分段器""" def test_merge_short_segments(self, short_segments): """Test 1: 连续短句被合并为段落""" from kb_workflow.srt.segmenter import SemanticSegmenter segmenter = SemanticSegmenter() paragraphs = segmenter.segment(short_segments) # 所有短句应该被合并 assert len(paragraphs) >= 1 # 验证所有内容都被保留 merged_text = ' '.join(paragraphs) assert "第一句短句" in merged_text assert "第二句短句" in merged_text assert "第三句短句" in merged_text def test_segment_on_topic_change(self, segments_with_gap): """Test 2: 主题变化时自动分段（基于时间间隔）""" from kb_workflow.srt.segmenter import SemanticSegmenter segmenter = SemanticSegmenter(gap_threshold=3.0) paragraphs = segmenter.segment(segments_with_gap) # 应该分成两段（因为中间有5秒停顿） assert len(paragraphs) == 2 assert "第一段" in paragraphs[0] assert "第二段" in paragraphs[1] def test_max_paragraph_length(self, long_text_segment): """Test 3: 单个段落不超过 500 字符""" from kb_workflow.srt.segmenter import SemanticSegmenter segmenter = SemanticSegmenter(max_paragraph_chars=500) paragraphs = segmenter.segment(long_text_segment) # 每个段落不应该超过限制 for para in paragraphs: assert len(para) <= 520 # 允许小范围超出（最后一块） def test_empty_segments(self): """空分段列表返回空列表""" from kb_workflow.srt.segmenter import SemanticSegmenter segmenter = SemanticSegmenter() paragraphs = segmenter.segment([]) assert paragraphs == [] # ============================================================================= # Test 3: ASR Corrector # ============================================================================= class TestASRCorrector: """测试 ASR 纠错器""" def test_correct_obvious_errors(self): """Test 1: 明显的同音字错误被纠正""" from kb_workflow.srt.corrector import ASRCorrector # 创建一个 mock AI 客户端 class MockAIClient: def generate(self, prompt: str, max_tokens: int = 16000) -> str: # 模拟纠错结果 if "什么" in prompt: return "什么是机器学习" return prompt corrector = ASRCorrector(ai_client=MockAIClient()) result = corrector.correct("什么是机器学习") # 纠错后的文本应该被返回 assert result is not None assert len(result) > 0 def test_correct_contextual_errors(self): """Test 2: 上下文相关的错误被纠正""" from kb_workflow.srt.corrector import ASRCorrector class MockAIClient: def generate(self, prompt: str, max_tokens: int = 16000) -> str: # 模拟纠错结果 return "这是一个测试文本，已经被纠正了。" corrector = ASRCorrector(ai_client=MockAIClient()) result = corrector.correct("这是一个测试文本，已经被纠正了。") assert "测试" in result def test_no_overcorrection(self): """Test 3: 正确文本不被过度修改""" from kb_workflow.srt.corrector import ASRCorrector original_text = "这是一段完全正确的中文文本，不需要修改。" class MockAIClient: def generate(self, prompt: str, max_tokens: int = 16000) -> str: # 返回原文，表示没有错误 return original_text corrector = ASRCorrector(ai_client=MockAIClient()) result = corrector.correct(original_text) # 正确文本应该基本保持不变 assert result == original_text def test_ai_unavailable_returns_original(self): """AI 不可用时返回原文""" from kb_workflow.srt.corrector import ASRCorrector # 没有提供 AI 客户端 corrector = ASRCorrector(ai_client=None) text = "这是原始文本" result = corrector.correct(text) # 应该返回原文 assert result == text def test_batch_processing(self): """长文本被分批处理""" from kb_workflow.srt.corrector import ASRCorrector # 创建超长文本 long_text = "这是一段很长的文本。" * 1000 # 约 10,000 字符 call_count = 0 class MockAIClient: def generate(self, prompt: str, max_tokens: int = 16000) -> str: nonlocal call_count call_count += 1 # 返回提示词中的文本部分 return "处理后的文本" corrector = ASRCorrector(ai_client=MockAIClient(), batch_size=5000) result = corrector.correct(long_text) # 应该被分批处理（至少调用2次） assert call_count >= 2 assert result is not None def test_empty_text(self): """空文本返回空字符串""" from kb_workflow.srt.corrector import ASRCorrector corrector = ASRCorrector(ai_client=None) result = corrector.correct("") assert result == "" # ============================================================================= # Test 4: SRTProcessor # ============================================================================= @pytest.fixture def srt_test_files(tmp_path: Path): """创建测试用的 SRT 和元数据文件""" # 创建 SRT 文件 srt_content = """1 00:00:00,000 --> 00:00:05,000 第一句字幕内容 2 00:00:05,000 --> 00:00:10,000 第二句字幕内容 3 00:00:10,000 --> 00:00:15,000 第三句字幕内容 """ srt_file = tmp_path / "test.srt" srt_file.write_text(srt_content, encoding="utf-8") # 创建元数据文件 metadata = { "title": "测试视频标题", "author": "测试作者", "platform": "bilibili", "bvid": "BV1test123", "duration": 60, "url": "https://www.bilibili.com/video/BV1test123", "pubtime": "2026-03-19" } metadata_file = tmp_path / "metadata.json" metadata_file.write_text(json.dumps(metadata, ensure_ascii=False), encoding="utf-8") # 输出文件路径 output_file = tmp_path / "output.md" return { "srt_file": str(srt_file), "metadata_file": str(metadata_file), "output_file": str(output_file) } class TestSRTProcessor: """测试 SRTProcessor 处理器""" def test_srt_processor_generates_markdown(self, srt_test_files): """SRTProcessor 能处理 SRT 文件生成 Markdown""" from kb_workflow.pipeline.processors.srt import SRTProcessor from kb_workflow.pipeline.schemas import SRTInput processor = SRTProcessor(enable_correction=False) input_data = SRTInput( content_id="BV1test123", platform="bilibili", srt_file=srt_test_files["srt_file"], metadata_file=srt_test_files["metadata_file"], output_file=srt_test_files["output_file"] ) # 运行处理 result = asyncio.run(processor.process(input_data)) assert result.success assert Path(result.output_file).exists() # 验证 Markdown 内容 content = Path(result.output_file).read_text(encoding="utf-8") assert "---" in content # frontmatter assert "title:" in content assert "测试视频标题" in content assert "## 内容" in content assert "第一句字幕内容" in content def test_srt_processor_with_yaml_frontmatter(self, srt_test_files): """输出包含 YAML frontmatter（元数据和标签）""" from kb_workflow.pipeline.processors.srt import SRTProcessor from kb_workflow.pipeline.schemas import SRTInput processor = SRTProcessor(enable_correction=False) input_data = SRTInput( content_id="BV1test123", platform="bilibili", srt_file=srt_test_files["srt_file"], metadata_file=srt_test_files["metadata_file"], output_file=srt_test_files["output_file"] ) result = asyncio.run(processor.process(input_data)) content = Path(result.output_file).read_text(encoding="utf-8") # 验证 frontmatter 字段 assert "title:" in content assert "date:" in content assert "author:" in content assert "platform:" in content assert "bvid:" in content assert "url:" in content def test_srt_processor_follows_base_processor_pattern(self): """处理器遵循 BaseProcessor 模式""" from kb_workflow.pipeline.processors.srt import SRTProcessor from kb_workflow.pipeline.processors.base import BaseProcessor # SRTProcessor 应该继承 BaseProcessor assert issubclass(SRTProcessor, BaseProcessor) # 应该有 STEP_NAME assert hasattr(SRTProcessor, 'STEP_NAME') assert SRTProcessor.STEP_NAME == "srt" def test_srt_processor_handles_missing_metadata(self, tmp_path: Path): """元数据文件缺失时能正常处理""" from kb_workflow.pipeline.processors.srt import SRTProcessor from kb_workflow.pipeline.schemas import SRTInput # 创建 SRT 文件 srt_content = "1\n00:00:00,000 --> 00:00:05,000\n测试内容\n" srt_file = tmp_path / "test.srt" srt_file.write_text(srt_content, encoding="utf-8") # 不存在的元数据文件 output_file = tmp_path / "output.md" processor = SRTProcessor(enable_correction=False) input_data = SRTInput( content_id="test123", platform="bilibili", srt_file=str(srt_file), metadata_file=str(tmp_path / "nonexistent.json"), output_file=str(output_file) ) result = asyncio.run(processor.process(input_data)) assert result.success assert Path(result.output_file).exists()