ReMe/tests4/unit/test_file_store.py
jinliyl 3cb2579ff7
refactor(steps): update auto-memory (#263)
* refactor(steps): update naming conventions in components and configuration

Updated naming conventions across multiple files, changing colon-separated names to underscore-separated format, and added new step definitions along with documentation updates.

Key changes:
- Replaced `Synchronizer` with `AutoMemory` as the counterpart component for cold-write operations
- Updated naming conventions in all related configuration files (e.g., `frontmatter:read` → `frontmatter_read`)
- Added new step definitions such as `submit_slug_updates` and `auto_memory`
- Updated relevant documentation
- Modified log output format for improved readability

* refactor(evolve): Refactor the auto-memory module and update related configurations

- Remove the old slug update commit step file
- Add new auto-memory planner and writer steps
- Update __init__.py to export the new step classes
- Modify the auto_memory configuration structure in default.yaml
- Update the slug field description for clearer explanation of its purpose

* up

* up

* refactor(tests): Move unit test directory from `tests4/unittest` to `tests4/unit`

Additionally, the assertion logic in test files has been updated: direct comparisons of `payload["notes"]` have been replaced with checks verifying the presence of paths and metadata within the response content. Furthermore, some test expectations have been simplified—for example, using `count` instead of asserting against specific note lists.

Specific changes include:
- Updating workflow configurations to align with the new test directory structure
- Modifying assertions across multiple test methods to make them more flexible and maintainable
- Cleaning up and optimizing parts of the test code structure

This is a comprehensive test refactoring effort aimed at improving test readability and robustness.

* Refactor(steps): Update memory writing logic and optimize JSON schema structure

Improved the write strategy description in `auto_memory_writer.yaml` to emphasize using `edit` over `write`.
Adjusted the `json_schema` structure in `base_step.py` to support the new function definition format.
Also corrected grammatical issues in the related documentation.

* Fix: Improve frontend data parsing error handling and update test files

Added capture and handling logic for YAML parsing exceptions, providing more detailed error messages when frontend data format issues occur. Also corrected the description text in a test file.
2026-05-29 12:07:44 +08:00

572 lines
20 KiB
Python

"""Tests for LocalFileStore and FaissLocalFileStore."""
# pylint: disable=protected-access
import asyncio
import hashlib
import importlib.util
import os
import tempfile
import warnings
import numpy as np
from reme4.components.embedding import BaseEmbeddingModel
from reme4.components.file_store import FaissLocalFileStore, LocalFileStore
from reme4.schema import FileChunk, FileNode
warnings.filterwarnings("ignore", category=DeprecationWarning, module="jieba")
warnings.filterwarnings("ignore", category=DeprecationWarning, module="pkg_resources")
_FAISS_AVAILABLE = importlib.util.find_spec("faiss") is not None
class FakeEmbeddingModel(BaseEmbeddingModel):
"""Deterministic stub: each lowercased word adds 1.0 at hash(word) % dim."""
def __init__(self, dimensions: int = 8, **kwargs):
super().__init__(model_name="fake", dimensions=dimensions, enable_cache=False, **kwargs)
async def _get_embeddings(self, input_text, **kwargs):
out = []
for t in input_text:
v = np.zeros(self.dimensions, dtype=np.float32)
for w in t.lower().split():
idx = int.from_bytes(hashlib.md5(w.encode()).digest()[:2], "big") % self.dimensions
v[idx] += 1.0
out.append(v.tolist())
return out
async def health_check(self, timeout: float = 2.0) -> bool:
self.is_healthy = True
return True
class temp_chdir:
"""Context manager to temporarily chdir into a path and restore on exit."""
def __init__(self, path):
self.path = path
self.old = None
def __enter__(self):
self.old = os.getcwd()
os.chdir(self.path)
return self
def __exit__(self, *exc):
os.chdir(self.old)
async def make_store(store_name: str = "test_store", **kwargs) -> LocalFileStore:
"""Build a started LocalFileStore with embedding disabled (no OpenAI dep)."""
store = LocalFileStore(name=store_name, embedding_model="", **kwargs)
await store.start()
return store
def make_file(
path: str,
text: str,
chunk_count: int = 1,
) -> tuple[FileNode, list[FileChunk]]:
"""Build a (FileNode, [FileChunk]) tuple ready for upsert_file."""
chunks = [
FileChunk(id=f"{path}::chunk{i}", path=path, text=f"{text} part{i}", start_line=i, end_line=i + 1)
for i in range(chunk_count)
]
node = FileNode(path=path, st_mtime=1.0, chunk_ids=[c.id for c in chunks])
return node, chunks
def test_upsert_single_file():
"""upsert_file with a one-element list stores chunks and node."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
node, chunks = make_file("a.md", "hello world", chunk_count=2)
await store.upsert([(node, chunks)])
# Chunks landed in memory
assert len(store.file_chunks) == 2
assert {c.path for c in store.file_chunks.values()} == {"a.md"}
# Node landed in graph
nodes = await store.get_nodes(["a.md"])
assert len(nodes) == 1
assert sorted(nodes[0].chunk_ids) == sorted([c.id for c in chunks])
await store.close()
print("✓ test_upsert_single_file passed")
asyncio.run(run())
def test_upsert_multiple_files():
"""upsert_file accepts a list of tuples and indexes them all."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
files = [make_file("a.md", "alpha"), make_file("b.md", "beta")]
await store.upsert(files)
assert len(store.file_chunks) == 2
paths = {n.path for n in await store.get_nodes()}
assert paths == {"a.md", "b.md"}
await store.close()
print("✓ test_upsert_multiple_files passed")
asyncio.run(run())
def test_upsert_replaces_old_chunks():
"""Re-upserting the same path points the node at the new chunk set."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
n1, c1 = make_file("a.md", "v1", chunk_count=2)
await store.upsert([(n1, c1)])
# Different chunks for the same path
n2 = FileNode(path="a.md", st_mtime=2.0)
c2 = [FileChunk(id="a.md::new", path="a.md", text="v2 only", start_line=0, end_line=1)]
n2.chunk_ids = [c.id for c in c2]
await store.upsert([(n2, c2)])
# The node now references the new chunk set, not the old one.
nodes = await store.get_nodes(["a.md"])
assert nodes[0].chunk_ids == ["a.md::new"]
assert "a.md::new" in store.file_chunks
await store.close()
print("✓ test_upsert_replaces_old_chunks passed")
asyncio.run(run())
def test_delete_by_path_single():
"""delete_by_path drops chunks and the node entry."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert([make_file("a.md", "alpha"), make_file("b.md", "beta")])
await store.delete("a.md")
assert all(c.path != "a.md" for c in store.file_chunks.values())
assert {n.path for n in await store.get_nodes()} == {"b.md"}
await store.close()
print("✓ test_delete_by_path_single passed")
asyncio.run(run())
def test_delete_by_path_list():
"""delete_by_path accepts a list of paths."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert(
[
make_file("a.md", "alpha"),
make_file("b.md", "beta"),
make_file("c.md", "gamma"),
],
)
await store.delete(["a.md", "b.md"])
assert {n.path for n in await store.get_nodes()} == {"c.md"}
assert all(c.path == "c.md" for c in store.file_chunks.values())
await store.close()
print("✓ test_delete_by_path_list passed")
asyncio.run(run())
def test_delete_by_path_missing_is_noop():
"""Deleting a nonexistent path is a no-op."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert([make_file("a.md", "alpha")])
before = len(store.file_chunks)
await store.delete("ghost.md")
assert len(store.file_chunks) == before
await store.close()
print("✓ test_delete_by_path_missing_is_noop passed")
asyncio.run(run())
def test_clear():
"""clear() empties chunks and the file graph."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert([make_file("a.md", "alpha"), make_file("b.md", "beta")])
await store.clear()
assert store.file_chunks == {}
assert await store.get_nodes() == []
await store.close()
print("✓ test_clear passed")
asyncio.run(run())
def test_keyword_search():
"""keyword_search returns matching chunks ranked by BM25 score."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert(
[
make_file("a.md", "python programming language"),
make_file("b.md", "java programming language"),
make_file("c.md", "python data analysis"),
],
)
results = await store.keyword_search("python", limit=5, search_filter={})
paths = {r.path for r in results}
assert "a.md" in paths or "c.md" in paths
# Each result should carry a keyword score.
for r in results:
assert r.scores.get("keyword", 0) > 0
await store.close()
print("✓ test_keyword_search passed")
asyncio.run(run())
def test_keyword_search_empty_query():
"""Empty/whitespace queries return no results."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert([make_file("a.md", "hello")])
assert await store.keyword_search("", limit=5, search_filter={}) == []
assert await store.keyword_search(" ", limit=5, search_filter={}) == []
await store.close()
print("✓ test_keyword_search_empty_query passed")
asyncio.run(run())
def test_vector_search_disabled_returns_empty():
"""Without an embedding model, vector_search returns []."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
await store.upsert([make_file("a.md", "hello")])
assert store.embedding_model is None
assert await store.vector_search("hello", limit=5, search_filter={}) == []
await store.close()
print("✓ test_vector_search_disabled_returns_empty passed")
asyncio.run(run())
def test_persistence_roundtrip():
"""close() dumps chunks; a fresh store loads them from disk."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
s1 = await make_store()
await s1.upsert([make_file("a.md", "alpha"), make_file("b.md", "beta")])
await s1.close()
s2 = await make_store()
assert {c.path for c in s2.file_chunks.values()} == {"a.md", "b.md"}
# Graph should also be persisted independently via its own dump.
assert {n.path for n in await s2.get_nodes()} == {"a.md", "b.md"}
await s2.close()
print("✓ test_persistence_roundtrip passed")
asyncio.run(run())
def test_rebuild_links_delegates_to_graph():
"""rebuild_links() on the store delegates to the underlying file_graph."""
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_store()
from reme4.schema import FileLink
node = FileNode(
path="a.md",
st_mtime=1.0,
links=[FileLink(source_path="a.md", target_path="b.md")],
)
chunks = [FileChunk(id="a::1", path="a.md", text="x", start_line=0, end_line=1)]
node.chunk_ids = [c.id for c in chunks]
await store.upsert([(node, chunks), make_file("b.md", "beta")])
await store.rebuild_links()
inlinks = await store.get_inlinks("b.md")
assert {lnk.source_path for lnk in inlinks} == {"a.md"}
await store.close()
print("✓ test_rebuild_links_delegates_to_graph passed")
asyncio.run(run())
# --- FaissLocalFileStore tests --------------------------------------------------
def _skip_if_no_faiss(name: str) -> bool:
if not _FAISS_AVAILABLE:
print(f"⊘ {name} skipped (faiss not installed)")
return True
return False
async def make_faiss_store(store_name: str = "test_faiss", **kwargs) -> FaissLocalFileStore:
"""Build a started FaissLocalFileStore wired to FakeEmbeddingModel (no API calls)."""
store = FaissLocalFileStore(name=store_name, embedding_model="fake", **kwargs)
fake = FakeEmbeddingModel()
# Replace the unresolved Dependency placeholder with a concrete instance and
# let start() cascade lifecycle to it via _owned.
store.embedding_model = fake
store._owned.append(fake)
await store.start()
return store
def test_faiss_vector_search_basic():
"""vector_search returns chunks ranked by cosine similarity to the query."""
if _skip_if_no_faiss("test_faiss_vector_search_basic"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_faiss_store()
await store.upsert(
[
make_file("a.md", "alpha"),
make_file("b.md", "beta"),
make_file("c.md", "alpha gamma"),
],
)
results = await store.vector_search("alpha", limit=3, search_filter={})
assert results, "vector_search returned no results"
for r in results:
assert "vector" in r.scores
assert r.scores["score"] == r.scores["vector"]
# Top hit should match an "alpha"-bearing doc.
assert results[0].path in {"a.md", "c.md"}
# All distinct chunks (no duplicates from tombstones).
assert len({r.id for r in results}) == len(results)
await store.close()
print("✓ test_faiss_vector_search_basic passed")
asyncio.run(run())
def test_faiss_persistence_roundtrip():
"""close() writes FAISS sidecar; a fresh store loads it without rebuilding."""
if _skip_if_no_faiss("test_faiss_persistence_roundtrip"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
s1 = await make_faiss_store()
await s1.upsert([make_file("a.md", "alpha"), make_file("b.md", "beta")])
r1 = await s1.vector_search("alpha", limit=2, search_filter={})
assert s1.faiss_path.exists() is False # not yet dumped
await s1.close()
assert s1.faiss_path.exists() and s1.faiss_idmap_path.exists()
s2 = await make_faiss_store()
assert s2._faiss_index is not None
assert s2._faiss_index.ntotal == 2
r2 = await s2.vector_search("alpha", limit=2, search_filter={})
assert [r.path for r in r2] == [r.path for r in r1]
await s2.close()
print("✓ test_faiss_persistence_roundtrip passed")
asyncio.run(run())
def test_faiss_delete_removes_from_search():
"""Deleting a file tombstones its chunks; subsequent search excludes them."""
if _skip_if_no_faiss("test_faiss_delete_removes_from_search"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_faiss_store()
await store.upsert([make_file("a.md", "alpha"), make_file("b.md", "alpha beta")])
await store.delete("a.md")
results = await store.vector_search("alpha", limit=5, search_filter={})
assert all(r.path != "a.md" for r in results)
# All tombstoned rows still present in id_map; live mapping shrank.
assert len(store._id_to_row) == 1
assert len(store._tombstones) == 1
await store.close()
print("✓ test_faiss_delete_removes_from_search passed")
asyncio.run(run())
def test_faiss_upsert_replaces_vectors():
"""Re-upserting a path with new chunk ids tombstones the old vectors."""
if _skip_if_no_faiss("test_faiss_upsert_replaces_vectors"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_faiss_store()
n1, c1 = make_file("a.md", "alpha", chunk_count=2)
await store.upsert([(n1, c1)])
assert store._faiss_index.ntotal == 2
assert len(store._tombstones) == 0
# New chunk ids for the same path → old ones become tombstones.
n2 = FileNode(path="a.md", st_mtime=2.0)
c2 = [FileChunk(id="a.md::new", path="a.md", text="gamma", start_line=0, end_line=1)]
n2.chunk_ids = [c.id for c in c2]
await store.upsert([(n2, c2)])
assert store._faiss_index.ntotal == 3 # 2 old + 1 new appended
assert len(store._tombstones) == 2 # both old rows tombstoned
assert "a.md::new" in store._id_to_row
results = await store.vector_search("gamma", limit=5, search_filter={})
assert results and results[0].id == "a.md::new"
await store.close()
print("✓ test_faiss_upsert_replaces_vectors passed")
asyncio.run(run())
def test_faiss_clear_empties_index():
"""clear() resets FAISS state and removes sidecar files."""
if _skip_if_no_faiss("test_faiss_clear_empties_index"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = await make_faiss_store()
await store.upsert([make_file("a.md", "alpha")])
await store.dump()
assert store.faiss_path.exists()
await store.clear()
assert store._faiss_index.ntotal == 0
assert store._id_map == [] and store._id_to_row == {}
assert not store.faiss_path.exists()
assert not store.faiss_idmap_path.exists()
assert await store.vector_search("alpha", limit=5, search_filter={}) == []
await store.close()
print("✓ test_faiss_clear_empties_index passed")
asyncio.run(run())
def test_faiss_rebuild_when_sidecar_missing():
"""If the FAISS sidecar is missing on load, the index rebuilds from chunks JSONL."""
if _skip_if_no_faiss("test_faiss_rebuild_when_sidecar_missing"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
s1 = await make_faiss_store()
await s1.upsert([make_file("a.md", "alpha"), make_file("b.md", "beta")])
await s1.close()
# Drop the FAISS sidecar but keep chunks JSONL — load() must rebuild.
s1.faiss_path.unlink()
s1.faiss_idmap_path.unlink()
s2 = await make_faiss_store()
assert s2._faiss_index is not None and s2._faiss_index.ntotal == 2
results = await s2.vector_search("alpha", limit=2, search_filter={})
assert any(r.path == "a.md" for r in results)
await s2.close()
print("✓ test_faiss_rebuild_when_sidecar_missing passed")
asyncio.run(run())
def test_faiss_disabled_without_embedding():
"""embedding_model="" → FAISS path stays dormant; vector_search returns []."""
if _skip_if_no_faiss("test_faiss_disabled_without_embedding"):
return
async def run():
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
store = FaissLocalFileStore(name="disabled", embedding_model="")
await store.start()
await store.upsert([make_file("a.md", "alpha")])
assert store.embedding_model is None
assert store._faiss_index is None
assert await store.vector_search("alpha", limit=5, search_filter={}) == []
await store.close()
print("✓ test_faiss_disabled_without_embedding passed")
asyncio.run(run())
if __name__ == "__main__":
print("\n=== LocalFileStore Tests ===")
test_upsert_single_file()
test_upsert_multiple_files()
test_upsert_replaces_old_chunks()
test_delete_by_path_single()
test_delete_by_path_list()
test_delete_by_path_missing_is_noop()
test_clear()
test_keyword_search()
test_keyword_search_empty_query()
test_vector_search_disabled_returns_empty()
test_persistence_roundtrip()
test_rebuild_links_delegates_to_graph()
print("\n=== FaissLocalFileStore Tests ===")
test_faiss_vector_search_basic()
test_faiss_persistence_roundtrip()
test_faiss_delete_removes_from_search()
test_faiss_upsert_replaces_vectors()
test_faiss_clear_empties_index()
test_faiss_rebuild_when_sidecar_missing()
test_faiss_disabled_without_embedding()
print("\n所有测试通过!")