refactor: route all chromadb access through ChromaBackend
Prerequisite for RFC 001 (plugin spec, #743). Removes every direct `import chromadb` outside the ChromaDB backend itself so the core modules depend only on the backend abstraction layer. Extends ChromaBackend with make_client, get_or_create_collection, delete_collection, create_collection, and backend_version. Adds update() to the BaseCollection contract. Non-backend callers (mcp_server, dedup, repair, migrate, cli) now go through the abstraction; tests patch ChromaBackend instead of chromadb. With this landed, the RFC 001 spec can be enforced and PalaceStore (#643) can ship as a plugin without touching core modules.
This commit is contained in:
+41
-65
@@ -412,12 +412,21 @@ def test_main_compress_dispatches():
|
||||
# ── cmd_repair ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _mock_backend_for(col=None, new_col=None):
|
||||
"""Build a mock ChromaBackend whose get_collection/create_collection return *col* / *new_col*."""
|
||||
mock_backend = MagicMock()
|
||||
if col is not None:
|
||||
mock_backend.get_collection.return_value = col
|
||||
if new_col is not None:
|
||||
mock_backend.create_collection.return_value = new_col
|
||||
return mock_backend
|
||||
|
||||
|
||||
@patch("mempalace.cli.MempalaceConfig")
|
||||
def test_cmd_repair_no_palace(mock_config_cls, tmp_path, capsys):
|
||||
mock_config_cls.return_value.palace_path = str(tmp_path / "nonexistent")
|
||||
args = argparse.Namespace(palace=None)
|
||||
mock_chromadb = MagicMock()
|
||||
with patch.dict("sys.modules", {"chromadb": mock_chromadb}):
|
||||
with patch("mempalace.backends.chroma.ChromaBackend"):
|
||||
cmd_repair(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "No palace found" in out
|
||||
@@ -429,8 +438,7 @@ def test_cmd_repair_requires_palace_database(mock_config_cls, tmp_path, capsys):
|
||||
palace_dir.mkdir()
|
||||
mock_config_cls.return_value.palace_path = str(palace_dir)
|
||||
args = argparse.Namespace(palace=None)
|
||||
mock_chromadb = MagicMock()
|
||||
with patch.dict("sys.modules", {"chromadb": mock_chromadb}):
|
||||
with patch("mempalace.backends.chroma.ChromaBackend"):
|
||||
cmd_repair(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "No palace database found" in out
|
||||
@@ -443,11 +451,9 @@ def test_cmd_repair_error_reading(mock_config_cls, tmp_path, capsys):
|
||||
(palace_dir / "chroma.sqlite3").write_text("db")
|
||||
mock_config_cls.return_value.palace_path = str(palace_dir)
|
||||
args = argparse.Namespace(palace=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.side_effect = Exception("corrupt db")
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
with patch.dict("sys.modules", {"chromadb": mock_chromadb}):
|
||||
mock_backend = MagicMock()
|
||||
mock_backend.get_collection.side_effect = Exception("corrupt db")
|
||||
with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend):
|
||||
cmd_repair(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "Error reading palace" in out
|
||||
@@ -460,13 +466,10 @@ def test_cmd_repair_zero_drawers(mock_config_cls, tmp_path, capsys):
|
||||
(palace_dir / "chroma.sqlite3").write_text("db")
|
||||
mock_config_cls.return_value.palace_path = str(palace_dir)
|
||||
args = argparse.Namespace(palace=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 0
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
with patch.dict("sys.modules", {"chromadb": mock_chromadb}):
|
||||
mock_backend = _mock_backend_for(col=mock_col)
|
||||
with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend):
|
||||
cmd_repair(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "Nothing to repair" in out
|
||||
@@ -479,7 +482,6 @@ def test_cmd_repair_success(mock_config_cls, tmp_path, capsys):
|
||||
(palace_dir / "chroma.sqlite3").write_text("db")
|
||||
mock_config_cls.return_value.palace_path = str(palace_dir)
|
||||
args = argparse.Namespace(palace=None, yes=True)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 2
|
||||
mock_col.get.return_value = {
|
||||
@@ -487,12 +489,9 @@ def test_cmd_repair_success(mock_config_cls, tmp_path, capsys):
|
||||
"documents": ["doc1", "doc2"],
|
||||
"metadatas": [{"wing": "a"}, {"wing": "b"}],
|
||||
}
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_new_col = MagicMock()
|
||||
mock_client.create_collection.return_value = mock_new_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
with patch.dict("sys.modules", {"chromadb": mock_chromadb}):
|
||||
mock_backend = _mock_backend_for(col=mock_col, new_col=mock_new_col)
|
||||
with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend):
|
||||
cmd_repair(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "Repair complete" in out
|
||||
@@ -506,20 +505,17 @@ def test_cmd_repair_aborts_without_confirmation(mock_config_cls, tmp_path, capsy
|
||||
(palace_dir / "chroma.sqlite3").write_text("db")
|
||||
mock_config_cls.return_value.palace_path = str(palace_dir)
|
||||
args = argparse.Namespace(palace=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 1
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
mock_backend = _mock_backend_for(col=mock_col)
|
||||
with (
|
||||
patch.dict("sys.modules", {"chromadb": mock_chromadb}),
|
||||
patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend),
|
||||
patch("builtins.input", return_value="n"),
|
||||
):
|
||||
cmd_repair(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "Aborted." in out
|
||||
mock_client.create_collection.assert_not_called()
|
||||
mock_backend.create_collection.assert_not_called()
|
||||
|
||||
|
||||
# ── cmd_compress ───────────────────────────────────────────────────────
|
||||
@@ -529,10 +525,10 @@ def test_cmd_repair_aborts_without_confirmation(mock_config_cls, tmp_path, capsy
|
||||
def test_cmd_compress_no_palace(mock_config_cls, capsys):
|
||||
mock_config_cls.return_value.palace_path = "/fake/palace"
|
||||
args = argparse.Namespace(palace=None, wing=None, dry_run=False, config=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_chromadb.PersistentClient.side_effect = Exception("no palace")
|
||||
mock_backend = MagicMock()
|
||||
mock_backend.get_collection.side_effect = Exception("no palace")
|
||||
with (
|
||||
patch.dict("sys.modules", {"chromadb": mock_chromadb}),
|
||||
patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend),
|
||||
pytest.raises(SystemExit),
|
||||
):
|
||||
cmd_compress(args)
|
||||
@@ -542,13 +538,10 @@ def test_cmd_compress_no_palace(mock_config_cls, capsys):
|
||||
def test_cmd_compress_no_drawers(mock_config_cls, capsys):
|
||||
mock_config_cls.return_value.palace_path = "/fake/palace"
|
||||
args = argparse.Namespace(palace=None, wing="mywing", dry_run=False, config=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.get.return_value = {"documents": [], "metadatas": [], "ids": []}
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
with patch.dict("sys.modules", {"chromadb": mock_chromadb}):
|
||||
mock_backend = _mock_backend_for(col=mock_col)
|
||||
with patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend):
|
||||
cmd_compress(args)
|
||||
out = capsys.readouterr().out
|
||||
assert "No drawers found" in out
|
||||
@@ -567,7 +560,6 @@ def _make_mock_dialect_module(dialect_instance):
|
||||
def test_cmd_compress_dry_run(mock_config_cls, capsys):
|
||||
mock_config_cls.return_value.palace_path = "/fake/palace"
|
||||
args = argparse.Namespace(palace=None, wing=None, dry_run=True, config=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.get.side_effect = [
|
||||
{
|
||||
@@ -577,9 +569,7 @@ def test_cmd_compress_dry_run(mock_config_cls, capsys):
|
||||
},
|
||||
{"documents": [], "metadatas": [], "ids": []},
|
||||
]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
mock_backend = _mock_backend_for(col=mock_col)
|
||||
|
||||
mock_dialect = MagicMock()
|
||||
mock_dialect.compress.return_value = "compressed"
|
||||
@@ -593,12 +583,9 @@ def test_cmd_compress_dry_run(mock_config_cls, capsys):
|
||||
}
|
||||
mock_dialect_mod = _make_mock_dialect_module(mock_dialect)
|
||||
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"chromadb": mock_chromadb,
|
||||
"mempalace.dialect": mock_dialect_mod,
|
||||
},
|
||||
with (
|
||||
patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend),
|
||||
patch.dict("sys.modules", {"mempalace.dialect": mock_dialect_mod}),
|
||||
):
|
||||
cmd_compress(args)
|
||||
out = capsys.readouterr().out
|
||||
@@ -613,22 +600,16 @@ def test_cmd_compress_with_config(mock_config_cls, tmp_path, capsys):
|
||||
config_file = tmp_path / "entities.json"
|
||||
config_file.write_text('{"people": [], "projects": []}')
|
||||
args = argparse.Namespace(palace=None, wing=None, dry_run=True, config=str(config_file))
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.get.return_value = {"documents": [], "metadatas": [], "ids": []}
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
mock_backend = _mock_backend_for(col=mock_col)
|
||||
|
||||
mock_dialect = MagicMock()
|
||||
mock_dialect_mod = _make_mock_dialect_module(mock_dialect)
|
||||
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"chromadb": mock_chromadb,
|
||||
"mempalace.dialect": mock_dialect_mod,
|
||||
},
|
||||
with (
|
||||
patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend),
|
||||
patch.dict("sys.modules", {"mempalace.dialect": mock_dialect_mod}),
|
||||
):
|
||||
cmd_compress(args)
|
||||
out = capsys.readouterr().out
|
||||
@@ -640,7 +621,6 @@ def test_cmd_compress_stores_results(mock_config_cls, capsys):
|
||||
"""Non-dry-run compress stores to mempalace_compressed collection."""
|
||||
mock_config_cls.return_value.palace_path = "/fake/palace"
|
||||
args = argparse.Namespace(palace=None, wing=None, dry_run=False, config=None)
|
||||
mock_chromadb = MagicMock()
|
||||
mock_col = MagicMock()
|
||||
mock_col.get.side_effect = [
|
||||
{
|
||||
@@ -650,11 +630,10 @@ def test_cmd_compress_stores_results(mock_config_cls, capsys):
|
||||
},
|
||||
{"documents": [], "metadatas": [], "ids": []},
|
||||
]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_comp_col = MagicMock()
|
||||
mock_client.get_or_create_collection.return_value = mock_comp_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
mock_backend = MagicMock()
|
||||
mock_backend.get_collection.return_value = mock_col
|
||||
mock_backend.get_or_create_collection.return_value = mock_comp_col
|
||||
|
||||
mock_dialect = MagicMock()
|
||||
mock_dialect.compress.return_value = "compressed"
|
||||
@@ -668,12 +647,9 @@ def test_cmd_compress_stores_results(mock_config_cls, capsys):
|
||||
}
|
||||
mock_dialect_mod = _make_mock_dialect_module(mock_dialect)
|
||||
|
||||
with patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"chromadb": mock_chromadb,
|
||||
"mempalace.dialect": mock_dialect_mod,
|
||||
},
|
||||
with (
|
||||
patch("mempalace.backends.chroma.ChromaBackend", return_value=mock_backend),
|
||||
patch.dict("sys.modules", {"mempalace.dialect": mock_dialect_mod}),
|
||||
):
|
||||
cmd_compress(args)
|
||||
out = capsys.readouterr().out
|
||||
|
||||
+19
-20
@@ -198,8 +198,15 @@ def test_dedup_source_group_query_failure_keeps():
|
||||
# ── show_stats ────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@patch("mempalace.dedup.chromadb")
|
||||
def test_show_stats(mock_chromadb, tmp_path):
|
||||
def _install_mock_backend(mock_backend_cls, collection):
|
||||
mock_backend = MagicMock()
|
||||
mock_backend.get_collection.return_value = collection
|
||||
mock_backend_cls.return_value = mock_backend
|
||||
return mock_backend
|
||||
|
||||
|
||||
@patch("mempalace.dedup.ChromaBackend")
|
||||
def test_show_stats(mock_backend_cls, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 5
|
||||
mock_col.get.side_effect = [
|
||||
@@ -215,9 +222,7 @@ def test_show_stats(mock_chromadb, tmp_path):
|
||||
},
|
||||
{"ids": []},
|
||||
]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
dedup.show_stats(palace_path=str(tmp_path)) # should not raise
|
||||
|
||||
@@ -227,13 +232,11 @@ def test_show_stats(mock_chromadb, tmp_path):
|
||||
|
||||
@patch("mempalace.dedup.dedup_source_group")
|
||||
@patch("mempalace.dedup.get_source_groups")
|
||||
@patch("mempalace.dedup.chromadb")
|
||||
def test_dedup_palace_dry_run(mock_chromadb, mock_groups, mock_dedup_group, tmp_path):
|
||||
@patch("mempalace.dedup.ChromaBackend")
|
||||
def test_dedup_palace_dry_run(mock_backend_cls, mock_groups, mock_dedup_group, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 10
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
mock_groups.return_value = {"a.txt": ["d1", "d2", "d3", "d4", "d5"]}
|
||||
mock_dedup_group.return_value = (["d1", "d2", "d3"], ["d4", "d5"])
|
||||
@@ -244,13 +247,11 @@ def test_dedup_palace_dry_run(mock_chromadb, mock_groups, mock_dedup_group, tmp_
|
||||
|
||||
@patch("mempalace.dedup.dedup_source_group")
|
||||
@patch("mempalace.dedup.get_source_groups")
|
||||
@patch("mempalace.dedup.chromadb")
|
||||
def test_dedup_palace_with_wing(mock_chromadb, mock_groups, mock_dedup_group, tmp_path):
|
||||
@patch("mempalace.dedup.ChromaBackend")
|
||||
def test_dedup_palace_with_wing(mock_backend_cls, mock_groups, mock_dedup_group, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 10
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
mock_groups.return_value = {}
|
||||
dedup.dedup_palace(palace_path=str(tmp_path), wing="test_wing", dry_run=True)
|
||||
@@ -259,13 +260,11 @@ def test_dedup_palace_with_wing(mock_chromadb, mock_groups, mock_dedup_group, tm
|
||||
|
||||
@patch("mempalace.dedup.dedup_source_group")
|
||||
@patch("mempalace.dedup.get_source_groups")
|
||||
@patch("mempalace.dedup.chromadb")
|
||||
def test_dedup_palace_no_groups(mock_chromadb, mock_groups, mock_dedup_group, tmp_path):
|
||||
@patch("mempalace.dedup.ChromaBackend")
|
||||
def test_dedup_palace_no_groups(mock_backend_cls, mock_groups, mock_dedup_group, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 3
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
mock_groups.return_value = {}
|
||||
dedup.dedup_palace(palace_path=str(tmp_path), dry_run=True)
|
||||
|
||||
+52
-62
@@ -66,22 +66,28 @@ def test_paginate_ids_offset_exception_fallback():
|
||||
# ── scan_palace ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_scan_palace_no_ids(mock_chromadb, tmp_path):
|
||||
def _install_mock_backend(mock_backend_cls, collection):
|
||||
"""Wire mock_backend_cls so ChromaBackend().get_collection(...) returns *collection*."""
|
||||
mock_backend = MagicMock()
|
||||
mock_backend.get_collection.return_value = collection
|
||||
mock_backend_cls.return_value = mock_backend
|
||||
return mock_backend
|
||||
|
||||
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_scan_palace_no_ids(mock_backend_cls, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 0
|
||||
mock_col.get.return_value = {"ids": []}
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
good, bad = repair.scan_palace(palace_path=str(tmp_path))
|
||||
assert good == set()
|
||||
assert bad == set()
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_scan_palace_all_good(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_scan_palace_all_good(mock_backend_cls, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 2
|
||||
# _paginate_ids call
|
||||
@@ -89,9 +95,7 @@ def test_scan_palace_all_good(mock_chromadb, tmp_path):
|
||||
{"ids": ["id1", "id2"]}, # paginate
|
||||
{"ids": ["id1", "id2"]}, # probe batch — both returned
|
||||
]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
good, bad = repair.scan_palace(palace_path=str(tmp_path))
|
||||
assert "id1" in good
|
||||
@@ -99,8 +103,8 @@ def test_scan_palace_all_good(mock_chromadb, tmp_path):
|
||||
assert len(bad) == 0
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_scan_palace_with_bad_ids(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_scan_palace_with_bad_ids(mock_backend_cls, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 2
|
||||
|
||||
@@ -117,26 +121,22 @@ def test_scan_palace_with_bad_ids(mock_chromadb, tmp_path):
|
||||
raise Exception("batch fail")
|
||||
|
||||
mock_col.get.side_effect = get_side_effect
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
good, bad = repair.scan_palace(palace_path=str(tmp_path))
|
||||
assert "good1" in good
|
||||
assert "bad1" in bad
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_scan_palace_with_wing_filter(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_scan_palace_with_wing_filter(mock_backend_cls, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 1
|
||||
mock_col.get.side_effect = [
|
||||
{"ids": ["id1"]}, # paginate
|
||||
{"ids": ["id1"]}, # probe
|
||||
]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
repair.scan_palace(palace_path=str(tmp_path), only_wing="test_wing")
|
||||
# Verify where filter was passed
|
||||
@@ -147,38 +147,36 @@ def test_scan_palace_with_wing_filter(mock_chromadb, tmp_path):
|
||||
# ── prune_corrupt ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_prune_corrupt_no_file(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_prune_corrupt_no_file(mock_backend_cls, tmp_path):
|
||||
# Should print message and return without error
|
||||
repair.prune_corrupt(palace_path=str(tmp_path))
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_prune_corrupt_dry_run(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_prune_corrupt_dry_run(mock_backend_cls, tmp_path):
|
||||
bad_file = tmp_path / "corrupt_ids.txt"
|
||||
bad_file.write_text("bad1\nbad2\n")
|
||||
repair.prune_corrupt(palace_path=str(tmp_path), confirm=False)
|
||||
# No chromadb calls in dry run
|
||||
mock_chromadb.PersistentClient.assert_not_called()
|
||||
# No backend calls in dry run
|
||||
mock_backend_cls.assert_not_called()
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_prune_corrupt_confirmed(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_prune_corrupt_confirmed(mock_backend_cls, tmp_path):
|
||||
bad_file = tmp_path / "corrupt_ids.txt"
|
||||
bad_file.write_text("bad1\nbad2\n")
|
||||
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.side_effect = [10, 8]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
repair.prune_corrupt(palace_path=str(tmp_path), confirm=True)
|
||||
mock_col.delete.assert_called_once()
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_prune_corrupt_delete_failure_fallback(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_prune_corrupt_delete_failure_fallback(mock_backend_cls, tmp_path):
|
||||
bad_file = tmp_path / "corrupt_ids.txt"
|
||||
bad_file.write_text("bad1\nbad2\n")
|
||||
|
||||
@@ -186,9 +184,7 @@ def test_prune_corrupt_delete_failure_fallback(mock_chromadb, tmp_path):
|
||||
mock_col.count.side_effect = [10, 8]
|
||||
# Batch delete fails, per-id succeeds
|
||||
mock_col.delete.side_effect = [Exception("batch fail"), None, None]
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
_install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
repair.prune_corrupt(palace_path=str(tmp_path), confirm=True)
|
||||
assert mock_col.delete.call_count == 3 # 1 batch + 2 individual
|
||||
@@ -197,29 +193,27 @@ def test_prune_corrupt_delete_failure_fallback(mock_chromadb, tmp_path):
|
||||
# ── rebuild_index ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_rebuild_index_no_palace(mock_chromadb, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_rebuild_index_no_palace(mock_backend_cls, tmp_path):
|
||||
nonexistent = str(tmp_path / "nope")
|
||||
repair.rebuild_index(palace_path=nonexistent)
|
||||
mock_chromadb.PersistentClient.assert_not_called()
|
||||
mock_backend_cls.assert_not_called()
|
||||
|
||||
|
||||
@patch("mempalace.repair.shutil")
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_rebuild_index_empty_palace(mock_chromadb, mock_shutil, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_rebuild_index_empty_palace(mock_backend_cls, mock_shutil, tmp_path):
|
||||
mock_col = MagicMock()
|
||||
mock_col.count.return_value = 0
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
||||
|
||||
repair.rebuild_index(palace_path=str(tmp_path))
|
||||
mock_client.delete_collection.assert_not_called()
|
||||
mock_backend.delete_collection.assert_not_called()
|
||||
|
||||
|
||||
@patch("mempalace.repair.shutil")
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_rebuild_index_success(mock_chromadb, mock_shutil, tmp_path):
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_rebuild_index_success(mock_backend_cls, mock_shutil, tmp_path):
|
||||
# Create a fake sqlite file
|
||||
sqlite_path = tmp_path / "chroma.sqlite3"
|
||||
sqlite_path.write_text("fake")
|
||||
@@ -233,10 +227,8 @@ def test_rebuild_index_success(mock_chromadb, mock_shutil, tmp_path):
|
||||
}
|
||||
|
||||
mock_new_col = MagicMock()
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.return_value = mock_col
|
||||
mock_client.create_collection.return_value = mock_new_col
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
mock_backend = _install_mock_backend(mock_backend_cls, mock_col)
|
||||
mock_backend.create_collection.return_value = mock_new_col
|
||||
|
||||
repair.rebuild_index(palace_path=str(tmp_path))
|
||||
|
||||
@@ -244,11 +236,9 @@ def test_rebuild_index_success(mock_chromadb, mock_shutil, tmp_path):
|
||||
mock_shutil.copy2.assert_called_once()
|
||||
assert "chroma.sqlite3" in str(mock_shutil.copy2.call_args)
|
||||
|
||||
# Verify: deleted and recreated with cosine
|
||||
mock_client.delete_collection.assert_called_once_with("mempalace_drawers")
|
||||
mock_client.create_collection.assert_called_once_with(
|
||||
"mempalace_drawers", metadata={"hnsw:space": "cosine"}
|
||||
)
|
||||
# Verify: deleted and recreated (cosine is the backend default)
|
||||
mock_backend.delete_collection.assert_called_once_with(str(tmp_path), "mempalace_drawers")
|
||||
mock_backend.create_collection.assert_called_once_with(str(tmp_path), "mempalace_drawers")
|
||||
|
||||
# Verify: used upsert not add
|
||||
mock_new_col.upsert.assert_called_once()
|
||||
@@ -256,11 +246,11 @@ def test_rebuild_index_success(mock_chromadb, mock_shutil, tmp_path):
|
||||
|
||||
|
||||
@patch("mempalace.repair.shutil")
|
||||
@patch("mempalace.repair.chromadb")
|
||||
def test_rebuild_index_error_reading(mock_chromadb, mock_shutil, tmp_path):
|
||||
mock_client = MagicMock()
|
||||
mock_client.get_collection.side_effect = Exception("corrupt")
|
||||
mock_chromadb.PersistentClient.return_value = mock_client
|
||||
@patch("mempalace.repair.ChromaBackend")
|
||||
def test_rebuild_index_error_reading(mock_backend_cls, mock_shutil, tmp_path):
|
||||
mock_backend = MagicMock()
|
||||
mock_backend.get_collection.side_effect = Exception("corrupt")
|
||||
mock_backend_cls.return_value = mock_backend
|
||||
|
||||
repair.rebuild_index(palace_path=str(tmp_path))
|
||||
mock_client.delete_collection.assert_not_called()
|
||||
mock_backend.delete_collection.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user