fix: Prevent Mock object issues in graph memory tests (#3627)
Co-authored-by: parshvadaftari <daftariparshva@gmail.com>
This commit is contained in:
co-authored by
parshvadaftari
parent
4a30745592
commit
eb2f8a3738
@@ -7,11 +7,25 @@ from mem0.memory.kuzu_memory import MemoryGraph
|
|||||||
class TestKuzu:
|
class TestKuzu:
|
||||||
"""Test that Kuzu memory works correctly"""
|
"""Test that Kuzu memory works correctly"""
|
||||||
|
|
||||||
|
# Create distinct embeddings that won't match with threshold=0.7
|
||||||
|
# Each embedding is mostly zeros with ones in different positions to ensure low similarity
|
||||||
|
alice_emb = np.zeros(384)
|
||||||
|
alice_emb[0:96] = 1.0
|
||||||
|
|
||||||
|
bob_emb = np.zeros(384)
|
||||||
|
bob_emb[96:192] = 1.0
|
||||||
|
|
||||||
|
charlie_emb = np.zeros(384)
|
||||||
|
charlie_emb[192:288] = 1.0
|
||||||
|
|
||||||
|
dave_emb = np.zeros(384)
|
||||||
|
dave_emb[288:384] = 1.0
|
||||||
|
|
||||||
embeddings = {
|
embeddings = {
|
||||||
"alice": np.random.uniform(0.0, 0.9, 384).tolist(),
|
"alice": alice_emb.tolist(),
|
||||||
"bob": np.random.uniform(0.0, 0.9, 384).tolist(),
|
"bob": bob_emb.tolist(),
|
||||||
"charlie": np.random.uniform(0.0, 0.9, 384).tolist(),
|
"charlie": charlie_emb.tolist(),
|
||||||
"dave": np.random.uniform(0.0, 0.9, 384).tolist(),
|
"dave": dave_emb.tolist(),
|
||||||
}
|
}
|
||||||
|
|
||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
@@ -26,6 +40,7 @@ class TestKuzu:
|
|||||||
|
|
||||||
# Mock graph store config
|
# Mock graph store config
|
||||||
config.graph_store.config.db = ":memory:"
|
config.graph_store.config.db = ":memory:"
|
||||||
|
config.graph_store.threshold = 0.7
|
||||||
|
|
||||||
# Mock LLM config
|
# Mock LLM config
|
||||||
config.llm.provider = "mock_llm"
|
config.llm.provider = "mock_llm"
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ class TestNeptuneMemory(unittest.TestCase):
|
|||||||
self.config = MagicMock()
|
self.config = MagicMock()
|
||||||
self.config.graph_store.config.endpoint = "neptune-graph://test-graph"
|
self.config.graph_store.config.endpoint = "neptune-graph://test-graph"
|
||||||
self.config.graph_store.config.base_label = True
|
self.config.graph_store.config.base_label = True
|
||||||
|
self.config.graph_store.threshold = 0.7
|
||||||
self.config.llm.provider = "openai_structured"
|
self.config.llm.provider = "openai_structured"
|
||||||
self.config.graph_store.llm = None
|
self.config.graph_store.llm = None
|
||||||
self.config.graph_store.custom_prompt = None
|
self.config.graph_store.custom_prompt = None
|
||||||
@@ -298,8 +299,8 @@ class TestNeptuneMemory(unittest.TestCase):
|
|||||||
|
|
||||||
# Verify the method calls
|
# Verify the method calls
|
||||||
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
|
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
|
||||||
self.memory_graph._search_source_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.9)
|
self.memory_graph._search_source_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
|
||||||
self.memory_graph._search_destination_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.9)
|
self.memory_graph._search_destination_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
|
||||||
self.memory_graph._add_entities_cypher.assert_called_once()
|
self.memory_graph._add_entities_cypher.assert_called_once()
|
||||||
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
|
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
|
||||||
|
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ class TestNeptuneMemory(unittest.TestCase):
|
|||||||
self.config = MagicMock()
|
self.config = MagicMock()
|
||||||
self.config.graph_store.config.endpoint = "neptune-db://test-graph"
|
self.config.graph_store.config.endpoint = "neptune-db://test-graph"
|
||||||
self.config.graph_store.config.base_label = True
|
self.config.graph_store.config.base_label = True
|
||||||
|
self.config.graph_store.threshold = 0.7
|
||||||
self.config.llm.provider = "openai_structured"
|
self.config.llm.provider = "openai_structured"
|
||||||
self.config.graph_store.llm = None
|
self.config.graph_store.llm = None
|
||||||
self.config.graph_store.custom_prompt = None
|
self.config.graph_store.custom_prompt = None
|
||||||
@@ -351,8 +352,8 @@ class TestNeptuneMemory(unittest.TestCase):
|
|||||||
|
|
||||||
# Verify the method calls
|
# Verify the method calls
|
||||||
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
|
self.assertEqual(self.mock_embedding_model.embed.call_count, 2)
|
||||||
self.memory_graph._search_source_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.9)
|
self.memory_graph._search_source_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
|
||||||
self.memory_graph._search_destination_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.9)
|
self.memory_graph._search_destination_node.assert_called_once_with(mock_embedding, self.user_id, threshold=0.7)
|
||||||
self.memory_graph._add_entities_cypher.assert_called_once()
|
self.memory_graph._add_entities_cypher.assert_called_once()
|
||||||
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
|
self.mock_graph.query.assert_called_once_with(mock_cypher, params=mock_params)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user