my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from unittest.mock import MagicMock, Mock, patch
#2
#3	# langchain_memgraph and rank_bm25 are optional deps — mock them so tests run without install
#4	_memgraph_mock = Mock()
#5	patch.dict("sys.modules", {
#6	"langchain_memgraph": _memgraph_mock,
#7	"langchain_memgraph.graphs": _memgraph_mock,
#8	"langchain_memgraph.graphs.memgraph": _memgraph_mock,
#9	"rank_bm25": Mock(),
#10	}).start()
#11
#12	from mem0.memory.memgraph_memory import MemoryGraph as MemgraphMemoryGraph # noqa: E402
#13	MemoryGraph = MemgraphMemoryGraph
#14
#15
#16	def _make_instance():
#17	with patch.object(MemoryGraph, "__init__", return_value=None):
#18	instance = MemoryGraph.__new__(MemoryGraph)
#19	instance.llm_provider = "openai"
#20	instance.llm = MagicMock()
#21	instance.embedding_model = MagicMock()
#22	instance.config = MagicMock()
#23	instance.config.graph_store.custom_prompt = None
#24	return instance
#25
#26
#27	class TestRetrieveNodesFromData:
#28	"""Tests for _retrieve_nodes_from_data in MemoryGraph."""
#29
#30	def test_normal_entities_extracted(self):
#31	instance = _make_instance()
#32	instance.llm.generate_response.return_value = {
#33	"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
#34	{"entity": "Alice", "entity_type": "person"},
#35	{"entity": "hiking", "entity_type": "activity"},
#36	]}}]
#37	}
#38	result = instance._retrieve_nodes_from_data("Alice loves hiking", {"user_id": "u1"})
#39	assert result == {"alice": "person", "hiking": "activity"}
#40
#41	def test_malformed_entity_missing_entity_type_is_skipped(self):
#42	"""LLM returns entity dict without entity_type — should skip it, keep valid ones.
#43	Reproduces the exact data from issue #4055."""
#44	instance = _make_instance()
#45	instance.llm.generate_response.return_value = {
#46	"tool_calls": [{"name": "extract_entities", "arguments": {"entities": [
#47	{"entity": "matrix multiplication", "entity_type": "task"},
#48	{"entity": "task"},
#49	{"entity": "ReLU", "entity_type": "task"},
#50	]}}]
#51	}
#52	result = instance._retrieve_nodes_from_data("some text", {"user_id": "u1"})
#53	assert "matrix_multiplication" in result
#54	assert "relu" in result
#55	assert "task" not in result
#56
#57	def test_none_tool_calls_returns_empty(self):
#58	instance = _make_instance()
#59	instance.llm.generate_response.return_value = {"tool_calls": None}
#60	result = instance._retrieve_nodes_from_data("hello world", {"user_id": "u1"})
#61	assert result == {}
#62
#63
#64	class TestEstablishNodesRelationsFromData:
#65	"""Tests for _establish_nodes_relations_from_data in MemoryGraph."""
#66
#67	def test_none_response_does_not_crash(self):
#68	"""openai_structured returns None when no relations found — must not crash.
#69	Exact crash from issue #4055: TypeError: 'NoneType' object is not subscriptable."""
#70	instance = _make_instance()
#71	instance.llm.generate_response.return_value = None
#72	result = instance._establish_nodes_relations_from_data(
#73	"Hello world", {"user_id": "u1"}, {}
#74	)
#75	assert result == []
#76
#77	def test_empty_tool_calls_returns_empty(self):
#78	instance = _make_instance()
#79	instance.llm.generate_response.return_value = {"tool_calls": []}
#80	result = instance._establish_nodes_relations_from_data(
#81	"Hello world", {"user_id": "u1"}, {}
#82	)
#83	assert result == []
#84
#85	def test_valid_entities_returned(self):
#86	instance = _make_instance()
#87	instance.llm.generate_response.return_value = {
#88	"tool_calls": [{"name": "add_entities", "arguments": {"entities": [
#89	{"source": "alice", "relationship": "loves", "destination": "hiking"}
#90	]}}]
#91	}
#92	result = instance._establish_nodes_relations_from_data(
#93	"Alice loves hiking", {"user_id": "u1"}, {"alice": "person"}
#94	)
#95	assert len(result) == 1
#96	assert result[0]["source"] == "alice"
#97

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public