diff --git a/tests/test_srt_processor.py b/tests/test_srt_processor.py index c178b62..cbdcadd 100644 --- a/tests/test_srt_processor.py +++ b/tests/test_srt_processor.py @@ -238,3 +238,105 @@ class TestSemanticSegmenter: 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 == ""