Tal Muskal
245 lines
7.7 KiB
Python
245 lines
7.7 KiB
Python
"""Tests for mempalace.palace_graph — graph traversal layer.
|
|
|
|
All ChromaDB access is mocked — no real database needed.
|
|
"""
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
|
|
def _make_fake_collection(metadatas, ids=None):
|
|
"""Create a mock collection that returns the given metadata in batches."""
|
|
if ids is None:
|
|
ids = [f"id_{i}" for i in range(len(metadatas))]
|
|
|
|
col = MagicMock()
|
|
col.count.return_value = len(metadatas)
|
|
|
|
def fake_get(limit=1000, offset=0, include=None):
|
|
batch_meta = metadatas[offset : offset + limit]
|
|
batch_ids = ids[offset : offset + limit]
|
|
return {"ids": batch_ids, "metadatas": batch_meta}
|
|
|
|
col.get.side_effect = fake_get
|
|
return col
|
|
|
|
|
|
# Patch chromadb at import time so palace_graph can be imported
|
|
with patch.dict("sys.modules", {"chromadb": MagicMock()}):
|
|
from mempalace.palace_graph import (
|
|
_fuzzy_match,
|
|
build_graph,
|
|
find_tunnels,
|
|
graph_stats,
|
|
traverse,
|
|
)
|
|
|
|
|
|
# --- build_graph ---
|
|
|
|
|
|
class TestBuildGraph:
|
|
def test_empty_collection(self):
|
|
col = _make_fake_collection([])
|
|
nodes, edges = build_graph(col=col)
|
|
assert nodes == {}
|
|
assert edges == []
|
|
|
|
def test_falsy_collection(self):
|
|
"""When col is explicitly falsy, build_graph returns empty."""
|
|
nodes, edges = build_graph(col=0)
|
|
assert nodes == {}
|
|
assert edges == []
|
|
|
|
def test_single_wing_no_edges(self):
|
|
col = _make_fake_collection(
|
|
[
|
|
{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-01"},
|
|
{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-02"},
|
|
]
|
|
)
|
|
nodes, edges = build_graph(col=col)
|
|
assert "auth" in nodes
|
|
assert nodes["auth"]["count"] == 2
|
|
assert edges == []
|
|
|
|
def test_multi_wing_creates_edges(self):
|
|
col = _make_fake_collection(
|
|
[
|
|
{
|
|
"room": "chromadb",
|
|
"wing": "wing_code",
|
|
"hall": "databases",
|
|
"date": "2026-01-01",
|
|
},
|
|
{
|
|
"room": "chromadb",
|
|
"wing": "wing_project",
|
|
"hall": "databases",
|
|
"date": "2026-01-02",
|
|
},
|
|
]
|
|
)
|
|
nodes, edges = build_graph(col=col)
|
|
assert "chromadb" in nodes
|
|
assert len(edges) == 1
|
|
assert edges[0]["wing_a"] == "wing_code"
|
|
assert edges[0]["wing_b"] == "wing_project"
|
|
assert edges[0]["hall"] == "databases"
|
|
|
|
def test_general_room_excluded(self):
|
|
col = _make_fake_collection(
|
|
[
|
|
{"room": "general", "wing": "wing_code", "hall": "misc", "date": ""},
|
|
]
|
|
)
|
|
nodes, edges = build_graph(col=col)
|
|
assert "general" not in nodes
|
|
|
|
def test_missing_wing_excluded(self):
|
|
col = _make_fake_collection(
|
|
[
|
|
{"room": "orphan", "wing": "", "hall": "misc", "date": ""},
|
|
]
|
|
)
|
|
nodes, edges = build_graph(col=col)
|
|
assert "orphan" not in nodes
|
|
|
|
def test_dates_capped_at_five(self):
|
|
col = _make_fake_collection(
|
|
[
|
|
{"room": "busy", "wing": "w", "hall": "h", "date": f"2026-01-{i:02d}"}
|
|
for i in range(1, 10)
|
|
]
|
|
)
|
|
nodes, _ = build_graph(col=col)
|
|
assert len(nodes["busy"]["dates"]) <= 5
|
|
|
|
|
|
# --- traverse ---
|
|
|
|
|
|
class TestTraverse:
|
|
def _build_col(self):
|
|
return _make_fake_collection(
|
|
[
|
|
{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-01"},
|
|
{"room": "login", "wing": "wing_code", "hall": "security", "date": "2026-01-01"},
|
|
{"room": "deploy", "wing": "wing_ops", "hall": "infra", "date": "2026-01-01"},
|
|
]
|
|
)
|
|
|
|
def test_traverse_known_room(self):
|
|
col = self._build_col()
|
|
result = traverse("auth", col=col)
|
|
assert isinstance(result, list)
|
|
rooms = [r["room"] for r in result]
|
|
assert "auth" in rooms
|
|
# login shares wing_code with auth
|
|
assert "login" in rooms
|
|
|
|
def test_traverse_unknown_room(self):
|
|
col = self._build_col()
|
|
result = traverse("nonexistent", col=col)
|
|
assert isinstance(result, dict)
|
|
assert "error" in result
|
|
assert "suggestions" in result
|
|
|
|
def test_traverse_max_hops(self):
|
|
col = self._build_col()
|
|
result = traverse("auth", col=col, max_hops=0)
|
|
# Only the start room itself at hop 0
|
|
assert len(result) == 1
|
|
assert result[0]["room"] == "auth"
|
|
|
|
|
|
# --- find_tunnels ---
|
|
|
|
|
|
class TestFindTunnels:
|
|
def _build_tunnel_col(self):
|
|
return _make_fake_collection(
|
|
[
|
|
{"room": "chromadb", "wing": "wing_code", "hall": "db", "date": "2026-01-01"},
|
|
{"room": "chromadb", "wing": "wing_project", "hall": "db", "date": "2026-01-02"},
|
|
{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-01"},
|
|
]
|
|
)
|
|
|
|
def test_find_all_tunnels(self):
|
|
col = self._build_tunnel_col()
|
|
tunnels = find_tunnels(col=col)
|
|
assert len(tunnels) == 1
|
|
assert tunnels[0]["room"] == "chromadb"
|
|
|
|
def test_find_tunnels_with_wing_filter(self):
|
|
col = self._build_tunnel_col()
|
|
tunnels = find_tunnels(wing_a="wing_code", col=col)
|
|
assert len(tunnels) == 1
|
|
|
|
def test_find_tunnels_no_match(self):
|
|
col = self._build_tunnel_col()
|
|
tunnels = find_tunnels(wing_a="wing_nonexistent", col=col)
|
|
assert tunnels == []
|
|
|
|
def test_find_tunnels_both_wings(self):
|
|
col = self._build_tunnel_col()
|
|
tunnels = find_tunnels(wing_a="wing_code", wing_b="wing_project", col=col)
|
|
assert len(tunnels) == 1
|
|
assert tunnels[0]["room"] == "chromadb"
|
|
|
|
|
|
# --- graph_stats ---
|
|
|
|
|
|
class TestGraphStats:
|
|
def test_empty_graph(self):
|
|
col = _make_fake_collection([])
|
|
stats = graph_stats(col=col)
|
|
assert stats["total_rooms"] == 0
|
|
assert stats["tunnel_rooms"] == 0
|
|
assert stats["total_edges"] == 0
|
|
|
|
def test_stats_with_data(self):
|
|
col = _make_fake_collection(
|
|
[
|
|
{"room": "chromadb", "wing": "wing_code", "hall": "db", "date": "2026-01-01"},
|
|
{"room": "chromadb", "wing": "wing_project", "hall": "db", "date": "2026-01-02"},
|
|
{"room": "auth", "wing": "wing_code", "hall": "security", "date": "2026-01-01"},
|
|
]
|
|
)
|
|
stats = graph_stats(col=col)
|
|
assert stats["total_rooms"] == 2
|
|
assert stats["tunnel_rooms"] == 1
|
|
assert stats["total_edges"] == 1
|
|
assert "wing_code" in stats["rooms_per_wing"]
|
|
|
|
|
|
# --- _fuzzy_match ---
|
|
|
|
|
|
class TestFuzzyMatch:
|
|
def test_exact_substring(self):
|
|
nodes = {"chromadb-setup": {}, "auth-module": {}, "deploy-config": {}}
|
|
result = _fuzzy_match("chromadb", nodes)
|
|
assert "chromadb-setup" in result
|
|
|
|
def test_partial_word_match(self):
|
|
nodes = {"chromadb-setup": {}, "auth-module": {}, "deploy-config": {}}
|
|
result = _fuzzy_match("auth", nodes)
|
|
assert "auth-module" in result
|
|
|
|
def test_no_match(self):
|
|
nodes = {"chromadb-setup": {}, "auth-module": {}}
|
|
result = _fuzzy_match("zzzzz", nodes)
|
|
assert result == []
|
|
|
|
def test_hyphenated_query(self):
|
|
nodes = {"riley-college-apps": {}, "college-prep": {}}
|
|
result = _fuzzy_match("riley-college", nodes)
|
|
assert "riley-college-apps" in result
|
|
|
|
def test_max_results(self):
|
|
nodes = {f"room-{i}": {} for i in range(20)}
|
|
result = _fuzzy_match("room", nodes, n=3)
|
|
assert len(result) <= 3
|