test(mcp): add multi-tenant and lazy-init tests for KG (#1136)

TestKGLazyCache covers the scenarios behind the lazy per-path refactor:

- test_lazy_init_no_import_side_effect: a fresh subprocess import does
  not create ~/.mempalace/knowledge_graph.sqlite3 (what closed PR #167
  was aiming at).
- test_get_kg_returns_same_instance: two _get_kg() calls under the same
  resolved path return the same object, cache has one entry.
- test_get_kg_different_paths_different_instances: rotating env var
  produces distinct KGs.
- test_multi_tenant_env_switch: the exact scenario from #1136 — write
  under path A, query under path B returns empty, switching back to A
  sees the fact.
- test_cache_thread_safe: 16 threads racing _get_kg() end up with one
  shared instance and one cache entry.
This commit is contained in:
mvalentsev
2026-04-24 12:53:19 +05:00
parent 9e730098e9
commit c69a622a18
+112
View File
@@ -8,6 +8,7 @@ via monkeypatch to avoid touching real data.
from datetime import datetime
import json
import os
import sys
import pytest
@@ -1143,3 +1144,114 @@ class TestCacheInvalidation:
for kwargs in captured["get"]:
assert "embedding_function" in kwargs
assert kwargs["embedding_function"] is not None
class TestKGLazyCache:
"""Lazy per-path KnowledgeGraph cache (issue #1136)."""
def test_lazy_init_no_import_side_effect(self, tmp_path):
"""Importing mcp_server must not create knowledge_graph.sqlite3.
Runs in a fresh subprocess with HOME pointed at tmp_path so the
assertion targets a clean filesystem, independent of conftest's
session-level HOME patch.
"""
import subprocess
import sys
kg_file = tmp_path / ".mempalace" / "knowledge_graph.sqlite3"
result = subprocess.run(
[sys.executable, "-c", "import mempalace.mcp_server"],
env={
"HOME": str(tmp_path),
"USERPROFILE": str(tmp_path),
"PATH": os.environ.get("PATH", ""),
"PYTHONPATH": os.environ.get("PYTHONPATH", ""),
},
capture_output=True,
text=True,
timeout=30,
)
assert result.returncode == 0, f"import failed: {result.stderr}"
assert not kg_file.exists(), f"import created sqlite file at {kg_file} as a side effect"
def test_get_kg_returns_same_instance(self, tmp_path, monkeypatch):
"""Two calls with the same resolved path return the same KG."""
from mempalace import mcp_server
monkeypatch.setattr(mcp_server, "_kg_by_path", {})
monkeypatch.setattr(mcp_server, "_palace_flag_given", True)
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_path))
kg1 = mcp_server._get_kg()
kg2 = mcp_server._get_kg()
assert kg1 is kg2
assert len(mcp_server._kg_by_path) == 1
def test_get_kg_different_paths_different_instances(self, tmp_path, monkeypatch):
"""Different palace paths map to different KG instances."""
from mempalace import mcp_server
tmp_a = tmp_path / "a"
tmp_b = tmp_path / "b"
tmp_a.mkdir()
tmp_b.mkdir()
monkeypatch.setattr(mcp_server, "_kg_by_path", {})
monkeypatch.setattr(mcp_server, "_palace_flag_given", True)
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_a))
kg_a = mcp_server._get_kg()
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_b))
kg_b = mcp_server._get_kg()
assert kg_a is not kg_b
assert len(mcp_server._kg_by_path) == 2
def test_multi_tenant_env_switch(self, tmp_path, monkeypatch):
"""The issue #1136 acceptance scenario.
Rotating MEMPALACE_PALACE_PATH between MCP tool calls must route
each call to the correct tenant's KG sqlite file.
"""
from mempalace import mcp_server
tmp_a = tmp_path / "tenant_a"
tmp_b = tmp_path / "tenant_b"
tmp_a.mkdir()
tmp_b.mkdir()
monkeypatch.setattr(mcp_server, "_kg_by_path", {})
monkeypatch.setattr(mcp_server, "_palace_flag_given", True)
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_a))
add_result = mcp_server.tool_kg_add(
subject="alice_secret",
predicate="owns",
object="repo_a",
)
assert add_result.get("success") is True, add_result
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_b))
query_b = mcp_server.tool_kg_query(entity="alice_secret")
assert query_b.get("count", 0) == 0, f"tenant B leaked tenant A's fact: {query_b}"
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_a))
query_a = mcp_server.tool_kg_query(entity="alice_secret")
assert query_a.get("count", 0) >= 1, f"tenant A lost its own fact: {query_a}"
def test_cache_thread_safe(self, tmp_path, monkeypatch):
"""Concurrent _get_kg() for the same path yields one instance."""
import concurrent.futures
from mempalace import mcp_server
monkeypatch.setattr(mcp_server, "_kg_by_path", {})
monkeypatch.setattr(mcp_server, "_palace_flag_given", True)
monkeypatch.setenv("MEMPALACE_PALACE_PATH", str(tmp_path))
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as pool:
results = list(pool.map(lambda _: mcp_server._get_kg(), range(16)))
ids = {id(kg) for kg in results}
assert len(ids) == 1, f"expected 1 unique instance, got {len(ids)}"
assert len(mcp_server._kg_by_path) == 1