157 lines
5.0 KiB
Python
157 lines
5.0 KiB
Python
from genesis.rag.embeddings import FakeEmbedder
|
||
from genesis.rag.store import RagStore
|
||
from genesis.rag.impact_rag import ImpactRAG
|
||
|
||
|
||
def test_retrieve_returns_relevant_chunk():
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
sources = [
|
||
("OrderController.java", "public class OrderController { 创建订单 }"),
|
||
("UserAuth.java", "public class UserAuth { 用户登录认证 }"),
|
||
]
|
||
rag.index("p1", sources)
|
||
res = rag.retrieve("p1", "OrderController 的影响范围", k=1)
|
||
assert res and "OrderController" in res[0]
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_reads_text_files(tmp_path):
|
||
(tmp_path / "src").mkdir()
|
||
(tmp_path / "src" / "OrderController.java").write_text(
|
||
"public class OrderController { 创建订单 }", encoding="utf-8"
|
||
)
|
||
(tmp_path / "src" / "binary.bin").write_bytes(b"\x00\x01")
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
n = rag.index_dir("p1", str(tmp_path / "src"))
|
||
assert n == 1
|
||
res = rag.retrieve("p1", "OrderController 的影响范围", k=1)
|
||
assert res and "OrderController" in res[0]
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_retrieve_empty_query_returns_empty():
|
||
# 空 query（仅空白）直接返回空列表，不调用向量检索
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
rag.index("p1", [("a.java", "创建订单")])
|
||
assert rag.retrieve("p1", " ", k=3) == []
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_empty_text_source_not_indexed():
|
||
# 空文本源：_split 回退为空片段，过滤后无 chunk
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
rag.index("p", [("empty.txt", "")])
|
||
assert rag.retrieve("p", "x", k=3) == []
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_non_dir_root_returns_zero():
|
||
# 非目录 root：is_dir() 为 False，返回 0
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
assert rag.index_dir("p", "不存在的路径") == 0
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_skips_empty_file(tmp_path):
|
||
# 空内容 .txt 文件不应被索引
|
||
(tmp_path / "empty.txt").write_text("", encoding="utf-8")
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
assert rag.index_dir("p", str(tmp_path)) == 0
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_skips_read_error(tmp_path, monkeypatch):
|
||
# 读取异常的文件被 except 跳过，整体不崩
|
||
(tmp_path / "a.txt").write_text("创建订单", encoding="utf-8")
|
||
monkeypatch.setattr(
|
||
__import__("pathlib").Path,
|
||
"read_text",
|
||
lambda *a, **k: (_ for _ in ()).throw(Exception),
|
||
)
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
assert rag.index_dir("p", str(tmp_path)) == 0
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_uppercase_ext_indexed(tmp_path):
|
||
# 大写扩展名（.JAVA）应经 .lower() 命中白名单并被索引
|
||
(tmp_path / "Foo.JAVA").write_text(
|
||
"public class Foo { 创建订单 }", encoding="utf-8"
|
||
)
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
n = rag.index_dir("p", str(tmp_path))
|
||
assert n == 1
|
||
res = rag.retrieve("p", "创建订单", k=1)
|
||
assert res and "Foo" in res[0]
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_empty_dir_returns_zero(tmp_path):
|
||
# 目录内无文本文件（sources 为空）返回 0
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
assert rag.index_dir("p", str(tmp_path)) == 0
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_dir_skips_nul_binary_text(tmp_path):
|
||
# 含 NUL 字节的 .txt 视为二进制误带扩展名，应跳过
|
||
(tmp_path / "bad.txt").write_bytes(b"hello\x00world")
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
assert rag.index_dir("p", str(tmp_path)) == 0
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_splits_long_single_paragraph():
|
||
# 单段落超长：触发 len(buf)+len(p) > size 且 buf 为空分支 -> 整体作为 1 个 chunk
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
rag.index("p", [("long.txt", "丙" * 900)])
|
||
chunks = rag.retrieve("p", "丙", k=5)
|
||
assert len(chunks) == 1
|
||
finally:
|
||
store.close()
|
||
|
||
|
||
def test_index_splits_multiple_paragraphs():
|
||
# 多段落（含空行）：触发拆行、空白行过滤、and buf 真分支 -> 2 个 chunk
|
||
store = RagStore(":memory:")
|
||
try:
|
||
rag = ImpactRAG(store, FakeEmbedder())
|
||
text = "甲" * 500 + "\n\n" + "乙" * 500
|
||
rag.index("p", [("multi.txt", text)])
|
||
chunks = rag.retrieve("p", "甲", k=5)
|
||
assert len(chunks) == 2
|
||
finally:
|
||
store.close()
|