my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from unittest.mock import MagicMock, patch
#2
#3	import pytest
#4
#5	from mem0 import Memory
#6	from mem0.configs.base import MemoryConfig
#7
#8
#9	class MockVectorMemory:
#10	"""Mock memory object for testing incomplete payloads."""
#11
#12	def __init__(self, memory_id: str, payload: dict, score: float = 0.8):
#13	self.id = memory_id
#14	self.payload = payload
#15	self.score = score
#16
#17
#18	@pytest.fixture
#19	def memory_client():
#20	with patch.object(Memory, "__init__", return_value=None):
#21	client = Memory()
#22	client.add = MagicMock(return_value={"results": [{"id": "1", "memory": "Name is John Doe.", "event": "ADD"}]})
#23	client.get = MagicMock(return_value={"id": "1", "memory": "Name is John Doe."})
#24	client.update = MagicMock(return_value={"message": "Memory updated successfully!"})
#25	client.delete = MagicMock(return_value={"message": "Memory deleted successfully!"})
#26	client.history = MagicMock(return_value=[{"memory": "I like Indian food."}, {"memory": "I like Italian food."}])
#27	client.get_all = MagicMock(return_value=["Name is John Doe.", "Name is John Doe. I like to code in Python."])
#28	yield client
#29
#30
#31	def test_create_memory(memory_client):
#32	data = "Name is John Doe."
#33	result = memory_client.add([{"role": "user", "content": data}], user_id="test_user")
#34	assert result["results"][0]["memory"] == data
#35
#36
#37	def test_get_memory(memory_client):
#38	data = "Name is John Doe."
#39	memory_client.add([{"role": "user", "content": data}], user_id="test_user")
#40	result = memory_client.get("1")
#41	assert result["memory"] == data
#42
#43
#44	def test_update_memory(memory_client):
#45	data = "Name is John Doe."
#46	memory_client.add([{"role": "user", "content": data}], user_id="test_user")
#47	new_data = "Name is John Kapoor."
#48	update_result = memory_client.update("1", text=new_data)
#49	assert update_result["message"] == "Memory updated successfully!"
#50
#51
#52	def test_delete_memory(memory_client):
#53	data = "Name is John Doe."
#54	memory_client.add([{"role": "user", "content": data}], user_id="test_user")
#55	delete_result = memory_client.delete("1")
#56	assert delete_result["message"] == "Memory deleted successfully!"
#57
#58
#59	def test_history(memory_client):
#60	data = "I like Indian food."
#61	memory_client.add([{"role": "user", "content": data}], user_id="test_user")
#62	memory_client.update("1", text="I like Italian food.")
#63	history = memory_client.history("1")
#64	assert history[0]["memory"] == "I like Indian food."
#65	assert history[1]["memory"] == "I like Italian food."
#66
#67
#68	def test_list_memories(memory_client):
#69	data1 = "Name is John Doe."
#70	data2 = "Name is John Doe. I like to code in Python."
#71	memory_client.add([{"role": "user", "content": data1}], user_id="test_user")
#72	memory_client.add([{"role": "user", "content": data2}], user_id="test_user")
#73	memories = memory_client.get_all(user_id="test_user")
#74	assert data1 in memories
#75	assert data2 in memories
#76
#77
#78	@patch('mem0.utils.factory.EmbedderFactory.create')
#79	@patch('mem0.utils.factory.VectorStoreFactory.create')
#80	@patch('mem0.utils.factory.LlmFactory.create')
#81	@patch('mem0.memory.storage.SQLiteManager')
#82	def test_collection_name_preserved_after_reset(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
#83	mock_embedder_factory.return_value = MagicMock()
#84	mock_vector_store = MagicMock()
#85	mock_vector_factory.return_value = mock_vector_store
#86	mock_llm_factory.return_value = MagicMock()
#87	mock_sqlite.return_value = MagicMock()
#88
#89	test_collection_name = "mem0"
#90	config = MemoryConfig()
#91	config.vector_store.config.collection_name = test_collection_name
#92
#93	memory = Memory(config)
#94
#95	assert memory.collection_name == test_collection_name
#96	assert memory.config.vector_store.config.collection_name == test_collection_name
#97
#98	memory.reset()
#99
#100	assert memory.collection_name == test_collection_name
#101	assert memory.config.vector_store.config.collection_name == test_collection_name
#102
#103	reset_calls = [call for call in mock_vector_factory.call_args_list if len(mock_vector_factory.call_args_list) > 2]
#104	if reset_calls:
#105	reset_config = reset_calls[-1][0][1]
#106	assert reset_config.collection_name == test_collection_name, f"Reset used wrong collection name: {reset_config.collection_name}"
#107
#108
#109	@patch('mem0.utils.factory.EmbedderFactory.create')
#110	@patch('mem0.utils.factory.VectorStoreFactory.create')
#111	@patch('mem0.utils.factory.LlmFactory.create')
#112	@patch('mem0.memory.storage.SQLiteManager')
#113	def test_search_handles_incomplete_payloads(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
#114	"""Test that search operations handle memory objects with missing 'data' key gracefully."""
#115	mock_embedder_factory.return_value = MagicMock()
#116	mock_vector_store = MagicMock()
#117	mock_vector_factory.return_value = mock_vector_store
#118	mock_llm_factory.return_value = MagicMock()
#119	mock_sqlite.return_value = MagicMock()
#120
#121	from mem0.memory.main import Memory as MemoryClass
#122	config = MemoryConfig()
#123	memory = MemoryClass(config)
#124
#125	# Create test data with both complete and incomplete payloads
#126	incomplete_memory = MockVectorMemory("mem_1", {"hash": "abc123"})
#127	complete_memory = MockVectorMemory("mem_2", {"data": "content", "hash": "def456"})
#128
#129	mock_vector_store.search.return_value = [incomplete_memory, complete_memory]
#130
#131	mock_embedder = MagicMock()
#132	mock_embedder.embed.return_value = [0.1, 0.2, 0.3]
#133	memory.embedding_model = mock_embedder
#134
#135	result = memory._search_vector_store("test", {"user_id": "test"}, 10)
#136
#137	assert len(result) == 2
#138	memories_by_id = {mem["id"]: mem for mem in result}
#139
#140	assert memories_by_id["mem_1"]["memory"] == ""
#141	assert memories_by_id["mem_2"]["memory"] == "content"
#142
#143
#144	@patch('mem0.utils.factory.EmbedderFactory.create')
#145	@patch('mem0.utils.factory.VectorStoreFactory.create')
#146	@patch('mem0.utils.factory.LlmFactory.create')
#147	@patch('mem0.memory.storage.SQLiteManager')
#148	def test_get_all_handles_nested_list_from_chroma(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
#149	"""
#150	Test that get_all() handles nested list return from Chroma/Milvus.
#151
#152	Issue #3674: Some vector stores return [[mem1, mem2]] instead of [mem1, mem2]
#153	This test ensures the unified unwrapping logic handles this correctly.
#154	"""
#155	mock_embedder_factory.return_value = MagicMock()
#156	mock_vector_store = MagicMock()
#157	mock_vector_factory.return_value = mock_vector_store
#158	mock_llm_factory.return_value = MagicMock()
#159	mock_sqlite.return_value = MagicMock()
#160
#161	from mem0.memory.main import Memory as MemoryClass
#162	config = MemoryConfig()
#163	memory = MemoryClass(config)
#164
#165	# Create test data
#166	mem1 = MockVectorMemory("mem_1", {"data": "My dog name is Sheru"})
#167	mem2 = MockVectorMemory("mem_2", {"data": "I like to code in Python"})
#168	mem3 = MockVectorMemory("mem_3", {"data": "I live in California"})
#169
#170	# Chroma/Milvus returns nested list: [[mem1, mem2, mem3]]
#171	mock_vector_store.list.return_value = [[mem1, mem2, mem3]]
#172
#173	result = memory._get_all_from_vector_store({"user_id": "test"}, 100)
#174
#175	# Should successfully unwrap and return 3 memories
#176	assert len(result) == 3
#177	assert result[0]["memory"] == "My dog name is Sheru"
#178	assert result[1]["memory"] == "I like to code in Python"
#179	assert result[2]["memory"] == "I live in California"
#180
#181
#182	@patch('mem0.utils.factory.EmbedderFactory.create')
#183	@patch('mem0.utils.factory.VectorStoreFactory.create')
#184	@patch('mem0.utils.factory.LlmFactory.create')
#185	@patch('mem0.memory.storage.SQLiteManager')
#186	def test_get_all_handles_tuple_from_qdrant(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
#187	"""
#188	Test that get_all() handles tuple return from Qdrant.
#189
#190	Qdrant returns: ([mem1, mem2], count)
#191	Should unwrap to [mem1, mem2]
#192	"""
#193	mock_embedder_factory.return_value = MagicMock()
#194	mock_vector_store = MagicMock()
#195	mock_vector_factory.return_value = mock_vector_store
#196	mock_llm_factory.return_value = MagicMock()
#197	mock_sqlite.return_value = MagicMock()
#198
#199	from mem0.memory.main import Memory as MemoryClass
#200	config = MemoryConfig()
#201	memory = MemoryClass(config)
#202
#203	mem1 = MockVectorMemory("mem_1", {"data": "Memory 1"})
#204	mem2 = MockVectorMemory("mem_2", {"data": "Memory 2"})
#205
#206	# Qdrant returns tuple: ([mem1, mem2], count)
#207	mock_vector_store.list.return_value = ([mem1, mem2], 100)
#208
#209	result = memory._get_all_from_vector_store({"user_id": "test"}, 100)
#210
#211	assert len(result) == 2
#212	assert result[0]["memory"] == "Memory 1"
#213	assert result[1]["memory"] == "Memory 2"
#214
#215
#216	@patch('mem0.utils.factory.EmbedderFactory.create')
#217	@patch('mem0.utils.factory.VectorStoreFactory.create')
#218	@patch('mem0.utils.factory.LlmFactory.create')
#219	@patch('mem0.memory.storage.SQLiteManager')
#220	def test_get_all_handles_flat_list_from_postgres(mock_sqlite, mock_llm_factory, mock_vector_factory, mock_embedder_factory):
#221	"""
#222	Test that get_all() handles flat list return from PostgreSQL.
#223
#224	PostgreSQL returns: [mem1, mem2]
#225	Should keep as-is without unwrapping
#226	"""
#227	mock_embedder_factory.return_value = MagicMock()
#228	mock_vector_store = MagicMock()
#229	mock_vector_factory.return_value = mock_vector_store
#230	mock_llm_factory.return_value = MagicMock()
#231	mock_sqlite.return_value = MagicMock()
#232
#233	from mem0.memory.main import Memory as MemoryClass
#234	config = MemoryConfig()
#235	memory = MemoryClass(config)
#236
#237	mem1 = MockVectorMemory("mem_1", {"data": "Memory 1"})
#238	mem2 = MockVectorMemory("mem_2", {"data": "Memory 2"})
#239
#240	# PostgreSQL returns flat list: [mem1, mem2]
#241	mock_vector_store.list.return_value = [mem1, mem2]
#242
#243	result = memory._get_all_from_vector_store({"user_id": "test"}, 100)
#244
#245	assert len(result) == 2
#246	assert result[0]["memory"] == "Memory 1"
#247	assert result[1]["memory"] == "Memory 2"
#248

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public