starmen1265
87b1919483
- Add SRTInput/SRTOutput schemas to pipeline - Create SRTProcessor following BaseProcessor pattern - Generate Markdown with YAML frontmatter - Support metadata extraction and formatting - Handle missing metadata gracefully Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
493 lines
16 KiB
Python
493 lines
16 KiB
Python
"""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()
|