Files
md-vector-db/tests/test_ingest.py
T
Serendipity 757bd8b336 feat: add Markdown splitter and document ingestor
- MarkdownSplitter: hybrid splitting by headings then paragraphs
- DocumentIngestor: read MD -> split -> embed -> store with dedup
- All 6 tests passing
2026-07-05 01:46:44 +08:00

88 lines
2.7 KiB
Python

"""文档入库测试."""
import tempfile
from pathlib import Path
import pytest
from src.core.ingest import MarkdownSplitter, DocumentIngestor
class TestMarkdownSplitter:
"""Markdown 分块器测试."""
@pytest.fixture
def splitter(self):
return MarkdownSplitter(max_size=1000, overlap=100)
def test_split_simple_document(self, splitter):
"""简单文档按标题拆分."""
md = """# 标题一
这是第一段内容。
## 标题二
这是第二段内容。
# 标题三
这是第三段内容。"""
chunks = splitter.split(md, source_file="test.md")
assert len(chunks) >= 3
# 每个 chunk 有元数据
for chunk in chunks:
assert "content" in chunk
assert chunk["source_file"] == "test.md"
def test_chunk_has_heading_metadata(self, splitter):
"""chunk 附带标题元数据."""
md = "# 配置指南\n这里是配置说明。"
chunks = splitter.split(md, source_file="config.md")
assert len(chunks) >= 1
title = chunks[0]["section_title"]
assert "配置指南" in title or title == ""
def test_long_section_is_split(self, splitter):
"""超长章节被进一步拆分."""
# 创建一个超过 max_size 的段落
long_text = "这是很长的文本。" * 300 # ~3000 字符
md = f"# 长章节\n{long_text}"
small_splitter = MarkdownSplitter(max_size=500, overlap=50)
chunks = small_splitter.split(md, source_file="long.md")
assert len(chunks) > 1
def test_empty_document(self, splitter):
"""空文档返回空列表."""
chunks = splitter.split("", source_file="empty.md")
assert chunks == []
def test_code_blocks_preserved(self, splitter):
"""代码块不被拆分."""
md = """# 代码示例
```python
def hello():
print("world")
```
"""
chunks = splitter.split(md, source_file="code.md")
assert len(chunks) >= 1
# 代码块内容应在某个 chunk 中
all_content = " ".join(c["content"] for c in chunks)
assert "def hello()" in all_content
class TestDocumentIngestor:
"""文档入库器测试."""
@pytest.fixture
def temp_md_dir(self):
with tempfile.TemporaryDirectory() as d:
# 创建测试 Markdown 文件
md_path = Path(d) / "test.md"
md_path.write_text("# 测试\n这是测试内容。", encoding="utf-8")
yield d
def test_read_markdown_file(self, temp_md_dir):
"""读取 Markdown 文件."""
content = Path(temp_md_dir + "/test.md").read_text(encoding="utf-8")
assert "测试" in content
assert "这是测试内容" in content