This repository has been archived on 2026-08-07. You can view files and clone it. You cannot open issues or pull requests or push a commit.
Files
kb-skills/tests/test_srt_processor.py
starmen1265 5983e95352 feat(15-02): integrate KBSettings config for semantic segmentation
- Modified SRTProcessor to read enable_semantic_segmentation from KBSettings
- Changed constructor parameter default from False to None (None = read from config)
- Added config priority: constructor parameter > KBSettings > default (True)
- Added logging for config status (INFO level)
- Added comprehensive tests for config integration (TestConfigIntegration)
- Added tests for output format compatibility (TestOutputFormatBenefits)
- Fixed MockAIClient to accept disable_reasoning parameter
- All 34 tests passing

Requirements satisfied:
- CONFIG-01: SRTProcessor reads from KBSettings
- CONFIG-02: Disabled config uses rule-based segmentation
- OUTPUT-01: LLM segmentation outputs H2 titles
- OUTPUT-02: title=null skips title output
- OUTPUT-03: Rule-based segmentation outputs "## 内容"

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-23 23:03:54 +08:00

936 lines
32 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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 os
import pytest
from pathlib import Path
from dataclasses import dataclass
from typing import List, Optional
from unittest.mock import Mock, patch
# =============================================================================
# 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, disable_reasoning: bool = False) -> 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,
enable_semantic_segmentation=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()
# =============================================================================
# Test 5: Semantic Segmentation Toggle (COST-03)
# =============================================================================
class TestSemanticSegmentationToggle:
"""测试语义分段配置开关 (COST-03)"""
def test_semantic_segmentation_enabled(self, sample_srt_file_for_toggle):
"""测试启用语义分段"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from unittest.mock import patch
with patch('kb_workflow.deep_analyzer.SimpleAIClient') as mock_client:
mock_client.return_value = Mock(api_key="test_key")
processor = SRTProcessor(
enable_semantic_segmentation=True
)
# 验证创建了 SemanticSegmentationEngine
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
assert isinstance(processor.segmenter, SemanticSegmentationEngine)
# 验证缓存已启用(默认)
assert processor.segmenter.strategy.cache is not None
def test_semantic_segmentation_disabled(self, sample_srt_file_for_toggle):
"""测试禁用语义分段(默认)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.srt.segmenter import SemanticSegmenter
processor = SRTProcessor(
enable_semantic_segmentation=False
)
# 验证使用规则分段器
assert isinstance(processor.segmenter, SemanticSegmenter)
def test_semantic_segmentation_with_cache(self, sample_srt_file_for_toggle):
"""测试启用语义分段和缓存"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from unittest.mock import patch
with patch('kb_workflow.deep_analyzer.SimpleAIClient') as mock_client:
mock_client.return_value = Mock(api_key="test_key")
processor = SRTProcessor(
enable_semantic_segmentation=True,
enable_cache=True
)
# 验证创建了 SemanticSegmentationEngine 且缓存已启用
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
assert isinstance(processor.segmenter, SemanticSegmentationEngine)
assert processor.segmenter.strategy.cache is not None
def test_semantic_segmentation_without_cache(self, sample_srt_file_for_toggle):
"""测试启用语义分段但禁用缓存"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from unittest.mock import patch
with patch('kb_workflow.deep_analyzer.SimpleAIClient') as mock_client:
mock_client.return_value = Mock(api_key="test_key")
processor = SRTProcessor(
enable_semantic_segmentation=True,
enable_cache=False
)
# 验证创建了 SemanticSegmentationEngine 但缓存未启用
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
assert isinstance(processor.segmenter, SemanticSegmentationEngine)
assert processor.segmenter.strategy.cache is None
@pytest.fixture
def sample_srt_file_for_toggle(self, tmp_path):
"""创建示例 SRT 文件"""
srt_content = """1
00:00:00,000 --> 00:00:01,000
第一句字幕
2
00:00:01,000 --> 00:00:02,000
第二句字幕
"""
srt_file = tmp_path / "test.srt"
srt_file.write_text(srt_content, encoding='utf-8')
return srt_file
# =============================================================================
# Test 6: Configuration Integration (CONFIG-01, CONFIG-02) - Plan 15-02
# =============================================================================
class TestConfigIntegration:
"""测试配置集成 (CONFIG-01, CONFIG-02, OUTPUT-01, OUTPUT-02, OUTPUT-03)"""
@pytest.fixture
def mock_env_base(self):
"""基础环境变量"""
return {
"KB_BASE_PATH": "/tmp/kb",
"KB_ARCHIVE_PATH": "/tmp/archive"
}
def test_srt_processor_reads_from_kbsettings(self, mock_env_base):
"""Test 1: enable_semantic_segmentation=None 时从 KBSettings 读取配置"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.config import get_settings
# 清除缓存以确保使用新的环境变量
get_settings.cache_clear()
with patch.dict(os.environ, {
**mock_env_base,
"KB_ENABLE_SEMANTIC_SEGMENTATION": "true"
}, clear=False):
get_settings.cache_clear()
processor = SRTProcessor()
# 验证从配置读取
assert processor.enable_semantic_segmentation is True
def test_srt_processor_constructor_overrides_config(self, mock_env_base):
"""Test 2: 显式传入参数时覆盖配置值(参数优先级高于配置)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.config import get_settings
get_settings.cache_clear()
with patch.dict(os.environ, {
**mock_env_base,
"KB_ENABLE_SEMANTIC_SEGMENTATION": "true"
}, clear=False):
get_settings.cache_clear()
# 显式传入 False 应该覆盖配置的 True
processor = SRTProcessor(enable_semantic_segmentation=False)
assert processor.enable_semantic_segmentation is False
def test_config_disabled_uses_rule_based(self, mock_env_base):
"""Test 3: 配置为 False 且无参数时使用规则分段"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.srt.segmenter import SemanticSegmenter
from kb_workflow.pipeline.config import get_settings
get_settings.cache_clear()
with patch.dict(os.environ, {
**mock_env_base,
"KB_ENABLE_SEMANTIC_SEGMENTATION": "false"
}, clear=False):
get_settings.cache_clear()
processor = SRTProcessor()
# 验证使用规则分段器
assert isinstance(processor.segmenter, SemanticSegmenter)
assert processor.enable_semantic_segmentation is False
def test_config_enabled_uses_llm_segmentation(self, mock_env_base):
"""Test 4: 配置为 True 且无参数时使用 LLM 分段"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.srt.semantic_engine import SemanticSegmentationEngine
from kb_workflow.pipeline.config import get_settings
get_settings.cache_clear()
with patch.dict(os.environ, {
**mock_env_base,
"KB_ENABLE_SEMANTIC_SEGMENTATION": "true"
}, clear=False):
get_settings.cache_clear()
with patch('kb_workflow.deep_analyzer.SimpleAIClient') as mock_client:
mock_client.return_value = Mock(api_key="test_key")
processor = SRTProcessor()
# 验证使用 LLM 分段引擎
assert isinstance(processor.segmenter, SemanticSegmentationEngine)
assert processor.enable_semantic_segmentation is True
def test_initialization_logs_config_status(self, mock_env_base, caplog):
"""Test 5: 初始化日志记录配置状态(INFO 级别)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.config import get_settings
import logging
get_settings.cache_clear()
with patch.dict(os.environ, {
**mock_env_base,
"KB_ENABLE_SEMANTIC_SEGMENTATION": "false"
}, clear=False):
get_settings.cache_clear()
with caplog.at_level(logging.INFO):
processor = SRTProcessor()
# 验证日志记录了配置状态
assert any("semantic_segmentation=" in record.message for record in caplog.records)
assert any("disabled" in record.message for record in caplog.records)
# =============================================================================
# Test 7: Output Format Compatibility (OUTPUT-01, OUTPUT-02, OUTPUT-03)
# =============================================================================
class TestOutputFormatCompatibility:
"""测试输出格式兼容性"""
@pytest.fixture
def srt_files_for_output(self, tmp_path):
"""创建输出格式测试用的文件"""
# 创建 SRT 文件
srt_content = """1
00:00:00,000 --> 00:00:05,000
这是第一段内容
2
00:00:05,000 --> 00:00:10,000
这是第二段内容
"""
srt_file = tmp_path / "test.srt"
srt_file.write_text(srt_content, encoding="utf-8")
# 创建元数据文件
metadata = {
"title": "输出格式测试",
"author": "测试作者",
"platform": "bilibili"
}
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)
}
def test_output_format_with_semantic_disabled(self, srt_files_for_output):
"""Test 3: 规则分段输出 "## 内容" 标题(OUTPUT-03)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
processor = SRTProcessor(
enable_semantic_segmentation=False,
enable_correction=False
)
input_data = SRTInput(
content_id="test123",
platform="bilibili",
srt_file=srt_files_for_output["srt_file"],
metadata_file=srt_files_for_output["metadata_file"],
output_file=srt_files_for_output["output_file"]
)
result = asyncio.run(processor.process(input_data))
# 验证输出文件
content = Path(result.output_file).read_text(encoding="utf-8")
# 验证 "## 内容" 标题(OUTPUT-03)
assert "## 内容" in content
# 验证段落内容直接输出
assert "这是第一段内容" in content
assert "这是第二段内容" in content
def test_generate_markdown_with_video_segments(self):
"""Test 1: LLM 分段输出带 H2 标题的 Markdown(OUTPUT-01)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.srt.semantic_engine import VideoSegment
processor = SRTProcessor(enable_correction=False)
# 创建 VideoSegment 对象模拟 LLM 分段结果
mock_segments = [
VideoSegment(title="第一段", content="这是第一段内容", start_index=0, end_index=10),
VideoSegment(title=None, content="这是第二段内容", start_index=10, end_index=20),
VideoSegment(title="第三段", content="这是第三段内容", start_index=20, end_index=30)
]
# 创建临时输出文件
import tempfile
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
output_file = Path(f.name)
try:
# 调用 _generate_markdown 方法
processor._generate_markdown(
output_file=output_file,
metadata={"title": "测试"},
paragraphs=mock_segments,
tags=[]
)
# 验证输出文件
content = output_file.read_text(encoding="utf-8")
# 验证 H2 标题(OUTPUT-01)
assert "## 第一段" in content
assert "## 第三段" in content
# 验证 title=None 时不输出标题(OUTPUT-02)
assert "## 第二段" not in content
assert "这是第二段内容" in content
# 验证内容正确输出
assert "这是第一段内容" in content
assert "这是第三段内容" in content
finally:
# 清理临时文件
if output_file.exists():
output_file.unlink()
def test_generate_markdown_with_plain_text(self):
"""Test 3b: 规则分段输出纯文本(带 "## 内容" 标题)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
processor = SRTProcessor(enable_correction=False)
# 创建纯文本列表模拟规则分段结果
mock_paragraphs = [
"这是第一段内容",
"这是第二段内容"
]
# 创建临时输出文件
import tempfile
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
output_file = Path(f.name)
try:
# 调用 _generate_markdown 方法
processor._generate_markdown(
output_file=output_file,
metadata={"title": "测试"},
paragraphs=mock_paragraphs,
tags=[]
)
# 验证输出文件
content = output_file.read_text(encoding="utf-8")
# 验证 "## 内容" 标题(OUTPUT-03)
assert "## 内容" in content
# 验证段落内容直接输出
assert "这是第一段内容" in content
assert "这是第二段内容" in content
finally:
# 清理临时文件
if output_file.exists():
output_file.unlink()
def test_generate_markdown_with_null_titles(self):
"""Test 2: 所有 title=None 时不输出分段标题(OUTPUT-02)"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.srt.semantic_engine import VideoSegment
processor = SRTProcessor(enable_correction=False)
# 创建全部 title=None 的 VideoSegment
mock_segments = [
VideoSegment(title=None, content="第一段内容", start_index=0, end_index=10),
VideoSegment(title=None, content="第二段内容", start_index=10, end_index=20)
]
# 创建临时输出文件
import tempfile
with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
output_file = Path(f.name)
try:
# 调用 _generate_markdown 方法
processor._generate_markdown(
output_file=output_file,
metadata={"title": "测试"},
paragraphs=mock_segments,
tags=[]
)
# 验证输出文件
content = output_file.read_text(encoding="utf-8")
# 验证没有分段标题(只有 frontmatter 和内容)
lines = content.split('\n')
# 找到 frontmatter 结束后的内容
in_content = False
for line in lines:
if line == "---":
if in_content:
break
in_content = True
continue
if in_content and line.startswith("## "):
# 只能有 "## 内容" 标题,不能有其他分段标题
assert line == "## 内容"
# 验证内容仍然存在
assert "第一段内容" in content
assert "第二段内容" in content
finally:
# 清理临时文件
if output_file.exists():
output_file.unlink()
def test_both_modes_valid_output(self, srt_files_for_output):
"""Test 4: 两种模式输出结构不同但都有效"""
from kb_workflow.pipeline.processors.srt import SRTProcessor
from kb_workflow.pipeline.schemas import SRTInput
# 测试规则分段模式
processor = SRTProcessor(
enable_semantic_segmentation=False,
enable_correction=False
)
input_data = SRTInput(
content_id="test123",
platform="bilibili",
srt_file=srt_files_for_output["srt_file"],
metadata_file=srt_files_for_output["metadata_file"],
output_file=srt_files_for_output["output_file"]
)
result = asyncio.run(processor.process(input_data))
assert result.success
assert Path(result.output_file).exists()
# 验证输出文件存在且非空
content = Path(result.output_file).read_text(encoding="utf-8")
assert len(content) > 0
assert "---" in content # frontmatter
assert "## 内容" in content