147 lines
4.9 KiB
Python
147 lines
4.9 KiB
Python
"""检索模块测试."""
|
|
import tempfile
|
|
|
|
import pytest
|
|
|
|
from src.core.config import EmbedConfig
|
|
from src.core.db import VectorDB
|
|
from src.core.embedder import create_embedder
|
|
from src.core.ingest import DocumentIngestor
|
|
from src.core.search import Searcher
|
|
|
|
|
|
@pytest.fixture
|
|
def searcher():
|
|
"""创建带测试数据的 Searcher."""
|
|
tmpdir = tempfile.mkdtemp()
|
|
db = VectorDB(persist_dir=tmpdir)
|
|
embedder = create_embedder(EmbedConfig(mode="local"))
|
|
ingestor = DocumentIngestor(db, embedder, "test_search")
|
|
|
|
# 入库一些测试文档
|
|
content = """# Python 入门
|
|
Python 是一种解释型编程语言。
|
|
|
|
## 安装 Python
|
|
从 python.org 下载安装包。
|
|
|
|
# 向量数据库
|
|
ChromaDB 是一个轻量级向量数据库。
|
|
|
|
## ChromaDB 安装
|
|
使用 pip install chromadb 安装。"""
|
|
ingestor.ingest_content(content, "guide.md")
|
|
|
|
return Searcher(db, embedder, "test_search")
|
|
|
|
|
|
class TestSearcher:
|
|
"""检索器测试."""
|
|
|
|
def test_search_returns_results(self, searcher):
|
|
"""搜索返回至少一条结果."""
|
|
results = searcher.search("Python 编程", top_k=3)
|
|
assert len(results) > 0
|
|
for r in results:
|
|
assert r["content"]
|
|
assert r["source_file"]
|
|
assert "score" in r
|
|
|
|
def test_search_scores_are_descending(self, searcher):
|
|
"""搜索结果按相似度降序排列."""
|
|
results = searcher.search("向量数据库", top_k=5)
|
|
scores = [r["score"] for r in results]
|
|
assert scores == sorted(scores, reverse=True)
|
|
|
|
def test_search_respects_top_k(self, searcher):
|
|
"""top_k 参数限制返回数量."""
|
|
results = searcher.search("安装", top_k=2)
|
|
assert len(results) <= 2
|
|
|
|
def test_search_returns_all_fields(self, searcher):
|
|
"""搜索结果包含完整字段."""
|
|
results = searcher.search("ChromaDB", top_k=1)
|
|
if results:
|
|
r = results[0]
|
|
assert "content" in r
|
|
assert "source_file" in r
|
|
assert "section_title" in r
|
|
assert "heading_level" in r
|
|
assert "chunk_index" in r
|
|
assert "score" in r
|
|
|
|
def test_search_no_results(self, searcher):
|
|
"""无语义匹配时不崩溃."""
|
|
results = searcher.search("xyzxyz不存在的内容abcabc", top_k=3)
|
|
assert isinstance(results, list)
|
|
|
|
def test_list_sources(self, searcher):
|
|
"""list_sources 返回已入库的源文件列表."""
|
|
sources = searcher.list_sources()
|
|
assert isinstance(sources, list)
|
|
|
|
def test_get_collection_info(self, searcher):
|
|
"""get_collection_info 返回 collection 信息."""
|
|
info = searcher.get_collection_info()
|
|
assert info["name"] == "test_search"
|
|
assert info["count"] > 0
|
|
|
|
def test_delete_by_source(self, searcher):
|
|
"""delete_by_source 删除源文件的所有 chunks."""
|
|
sources_before = searcher.list_sources()
|
|
if sources_before:
|
|
target = sources_before[0]
|
|
result = searcher.delete_by_source(target)
|
|
assert result is True
|
|
sources_after = searcher.list_sources()
|
|
assert target not in sources_after
|
|
|
|
def test_search_with_source_filter(self, searcher):
|
|
"""带 source_file 过滤的搜索."""
|
|
sources = searcher.list_sources()
|
|
if sources:
|
|
results = searcher.search("测试", top_k=3, source_file=sources[0])
|
|
assert isinstance(results, list)
|
|
for r in results:
|
|
assert r["source_file"] == sources[0]
|
|
|
|
def test_delete_by_source_nonexistent(self, searcher):
|
|
"""删除不存在的源文件返回 False."""
|
|
result = searcher.delete_by_source("nonexistent_file_xyz.md")
|
|
assert result is False
|
|
|
|
def test_list_sources_empty_collection(self, tmp_path):
|
|
"""空 collection 的 list_sources 返回空列表."""
|
|
from src.core.config import EmbedConfig
|
|
from src.core.db import VectorDB
|
|
from src.core.embedder import create_embedder
|
|
from src.core.search import Searcher
|
|
|
|
db = VectorDB(persist_dir=str(tmp_path))
|
|
embedder = create_embedder(EmbedConfig(mode="local"))
|
|
searcher = Searcher(db, embedder, "empty_coll")
|
|
sources = searcher.list_sources()
|
|
assert sources == []
|
|
|
|
|
|
def test_list_sources_empty():
|
|
"""空 collection 的 list_sources 返回空列表."""
|
|
from src.core.search import Searcher
|
|
|
|
class EmptyColl:
|
|
def count(self): return 0
|
|
def get(self, **kwargs):
|
|
return {"ids": [], "documents": [], "metadatas": []}
|
|
|
|
class EmptyDB:
|
|
def get_or_create_collection(self, name): return EmptyColl()
|
|
def list_collections(self): return []
|
|
|
|
class FakeEmb:
|
|
@property
|
|
def dimension(self): return 4
|
|
def embed(self, texts): return [[0.0] * 4]
|
|
|
|
s = Searcher(EmptyDB(), FakeEmb(), "empty")
|
|
assert s.list_sources() == []
|