File size: 5,038 Bytes
c3a3710 7c8b011 c3a3710 7c8b011 c3a3710 7c8b011 c3a3710 7c8b011 c3a3710 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import asyncio
import os
import shutil
import pytest
from pathlib import Path
import numpy as np
# Set dummy test config environment
os.environ["HAIM_API_KEY"] = "test-key"
from mnemocore.core.config import HAIMConfig, PathsConfig, TierConfig
from mnemocore.core.binary_hdv import TextEncoder, BinaryHDV
from mnemocore.core.hnsw_index import HNSWIndexManager
from mnemocore.core.engine import HAIMEngine
from mnemocore.core.tier_manager import TierManager
from unittest.mock import patch
@pytest.fixture(autouse=True)
def setup_test_env():
# Force all components to use test_data_verify as their data dir
# to prevent polluting/reading the user's real ./data folder
test_dir = Path("./test_data_verify")
test_dir.mkdir(exist_ok=True)
cfg = HAIMConfig(
paths=PathsConfig(
data_dir=str(test_dir),
warm_mmap_dir=str(test_dir / "warm"),
cold_archive_dir=str(test_dir / "cold")
)
)
with patch('mnemocore.core.config.get_config', return_value=cfg), \
patch('mnemocore.core.hnsw_index.get_config', return_value=cfg), \
patch('mnemocore.core.engine.get_config', return_value=cfg):
yield cfg
if test_dir.exists():
shutil.rmtree(test_dir)
@pytest.mark.asyncio
async def test_text_encoder_normalization():
"""Verify BUG-02: Text normalization fixes identical string variances"""
encoder = TextEncoder(dimension=16384)
hdv1 = encoder.encode("Hello World")
hdv2 = encoder.encode("hello, world!")
assert (hdv1.data == hdv2.data).all(), "Normalization failed: Different HDVs for identical texts"
def test_hnsw_singleton():
"""Verify BUG-08: HNSWIndexManager is a thread-safe singleton"""
HNSWIndexManager._instance = None
idx1 = HNSWIndexManager(dimension=16384)
idx2 = HNSWIndexManager(dimension=16384)
assert idx1 is idx2, "HNSWIndexManager is not a singleton"
def test_hnsw_index_add_search():
"""Verify BUG-01 & BUG-03: Vector cache lost / Position mapping"""
HNSWIndexManager._instance = None
idx = HNSWIndexManager(dimension=16384)
# Optional cleanup if it's reused
idx._id_map = []
idx._vector_store = []
if idx._index:
idx._index.reset()
vec1 = BinaryHDV.random(16384)
vec2 = BinaryHDV.random(16384)
idx.add("test_node_1", vec1.data)
idx.add("test_node_2", vec2.data)
assert "test_node_1" in idx._id_map, "ID Map does not contain node 1"
assert "test_node_2" in idx._id_map, "ID Map does not contain node 2"
# The search should return test_node_1 as the top result for vec1.data
res = idx.search(vec1.data, top_k=1)
assert res[0][0] == "test_node_1", f"Incorrect search return: {res}"
@pytest.mark.asyncio
async def test_agent_isolation():
"""Verify BUG-09: Agent namespace isolation via engine and tier manager"""
HNSWIndexManager._instance = None
test_data_dir = Path("./test_data_verify")
test_data_dir.mkdir(exist_ok=True)
config = HAIMConfig(
qdrant=None,
paths=PathsConfig(
data_dir=str(test_data_dir),
warm_mmap_dir=str(test_data_dir / "warm"),
cold_archive_dir=str(test_data_dir / "cold")
),
tiers_hot=TierConfig(max_memories=1000, ltp_threshold_min=0.0)
)
# Prevent newly created memories (LTP=0.5) from being eagerly demoted
# We run purely local/in-memory for this unit test
tier_manager = TierManager(config=config, qdrant_store=None)
engine = HAIMEngine(
persist_path=str(test_data_dir / "memory.jsonl"),
config=config,
tier_manager=tier_manager
)
try:
await engine.initialize()
# Store two memories, isolated
await engine.store("Secret logic for agent 1", metadata={"agent_id": "agent_alpha"})
await engine.store("Public logic for agent 2", metadata={"agent_id": "agent_beta"})
# Search global
res_global = await engine.query("logic", top_k=5)
# We expect 2 given we just pushed 2
assert len(res_global) >= 2, f"Global search should return at least 2 memories, got {len(res_global)}"
# Search isolated by agent_alpha
res_isolated = await engine.query("logic", top_k=5, metadata_filter={"agent_id": "agent_alpha"})
assert len(res_isolated) > 0, "Should find at least 1 memory for agent_alpha"
for nid, score in res_isolated:
node = await engine.get_memory(nid)
assert node.metadata.get("agent_id") == "agent_alpha", "Found leaked memory from another agent namespace!"
finally:
await engine.close()
# Clean up test dir
if test_data_dir.exists():
shutil.rmtree(test_data_dir)
if __name__ == "__main__":
pytest.main(["-v", __file__])
|