36 lines
1.2 KiB
Python
36 lines
1.2 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()
|