79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
"""Splitter 注册表和 TextSplitter 测试."""
|
|
from src.core.splitters import (
|
|
SUPPORTED_SUFFIXES,
|
|
MarkdownSplitter,
|
|
TextSplitter,
|
|
get_splitter,
|
|
register_splitter,
|
|
)
|
|
|
|
|
|
class TestTextSplitter:
|
|
"""TextSplitter 纯文本分块测试."""
|
|
|
|
def test_empty_text(self):
|
|
s = TextSplitter()
|
|
assert s.split("") == []
|
|
assert s.split(" \n\n ") == []
|
|
|
|
def test_short_text_single_chunk(self):
|
|
s = TextSplitter(max_size=1000)
|
|
chunks = s.split("这是一段短文本。", source_file="test.txt")
|
|
assert len(chunks) == 1
|
|
assert chunks[0]["source_file"] == "test.txt"
|
|
assert chunks[0]["content"] == "这是一段短文本。"
|
|
|
|
def test_long_paragraph_split(self):
|
|
s = TextSplitter(max_size=50, overlap=10)
|
|
long_text = "这是第一句。" * 20
|
|
chunks = s.split(long_text, source_file="long.txt")
|
|
assert len(chunks) > 1
|
|
for c in chunks:
|
|
assert len(c["content"]) <= 60 # max_size + 少许容差
|
|
|
|
def test_paragraph_boundary_split(self):
|
|
s = TextSplitter(max_size=100)
|
|
text = "短段落A。\n\n短段落B。\n\n短段落C。"
|
|
chunks = s.split(text)
|
|
assert len(chunks) >= 1
|
|
assert all("content" in c for c in chunks)
|
|
|
|
def test_chunk_metadata(self):
|
|
s = TextSplitter()
|
|
chunks = s.split("测试内容。", source_file="doc.txt")
|
|
assert chunks[0]["source_file"] == "doc.txt"
|
|
assert chunks[0]["section_title"] == ""
|
|
assert chunks[0]["heading_level"] == 0
|
|
assert chunks[0]["chunk_index"] == 0
|
|
|
|
|
|
class TestRegistry:
|
|
"""注册表测试."""
|
|
|
|
def test_get_splitter_for_md(self):
|
|
s = get_splitter("doc.md")
|
|
assert isinstance(s, MarkdownSplitter)
|
|
|
|
def test_get_splitter_for_txt(self):
|
|
s = get_splitter("notes.txt")
|
|
assert isinstance(s, TextSplitter)
|
|
|
|
def test_get_splitter_fallback(self):
|
|
s = get_splitter("data.xyz")
|
|
assert isinstance(s, TextSplitter)
|
|
|
|
def test_supported_suffixes(self):
|
|
assert ".md" in SUPPORTED_SUFFIXES
|
|
assert ".txt" in SUPPORTED_SUFFIXES
|
|
assert ".pdf" in SUPPORTED_SUFFIXES
|
|
assert ".html" in SUPPORTED_SUFFIXES
|
|
|
|
def test_custom_register(self):
|
|
class FakeSplitter:
|
|
def __init__(self, max_size=1000, overlap=100): pass
|
|
def split(self, text, source_file=""): return []
|
|
|
|
register_splitter(".fake", FakeSplitter)
|
|
s = get_splitter("test.fake")
|
|
assert isinstance(s, FakeSplitter)
|