diff --git a/tests/vector_stores/test_weaviate.py b/tests/vector_stores/test_weaviate.py index 96776e43c..680ec8455 100644 --- a/tests/vector_stores/test_weaviate.py +++ b/tests/vector_stores/test_weaviate.py @@ -1,220 +1,207 @@ -# import os -# import uuid -# import httpx -# import unittest -# from unittest.mock import MagicMock, patch +import os +import uuid +import httpx +import unittest +from unittest.mock import MagicMock, patch -# import dotenv -# import weaviate -# from weaviate.classes.query import MetadataQuery, Filter -# from weaviate.exceptions import UnexpectedStatusCodeException +import dotenv +import weaviate +from weaviate.exceptions import UnexpectedStatusCodeException -# from mem0.vector_stores.weaviate import Weaviate, OutputData +from mem0.vector_stores.weaviate import Weaviate -# class TestWeaviateDB(unittest.TestCase): -# @classmethod -# def setUpClass(cls): -# dotenv.load_dotenv() +class TestWeaviateDB(unittest.TestCase): + @classmethod + def setUpClass(cls): + dotenv.load_dotenv() -# cls.original_env = { -# 'WEAVIATE_CLUSTER_URL': os.getenv('WEAVIATE_CLUSTER_URL', 'http://localhost:8080'), -# 'WEAVIATE_API_KEY': os.getenv('WEAVIATE_API_KEY', 'test_api_key'), -# } + cls.original_env = { + "WEAVIATE_CLUSTER_URL": os.getenv("WEAVIATE_CLUSTER_URL", "http://localhost:8080"), + "WEAVIATE_API_KEY": os.getenv("WEAVIATE_API_KEY", "test_api_key"), + } -# os.environ['WEAVIATE_CLUSTER_URL'] = 'http://localhost:8080' -# os.environ['WEAVIATE_API_KEY'] = 'test_api_key' + os.environ["WEAVIATE_CLUSTER_URL"] = "http://localhost:8080" + os.environ["WEAVIATE_API_KEY"] = "test_api_key" -# def setUp(self): -# self.client_mock = MagicMock(spec=weaviate.WeaviateClient) -# self.client_mock.collections = MagicMock() -# self.client_mock.collections.exists.return_value = False -# self.client_mock.collections.create.return_value = None -# self.client_mock.collections.delete.return_value = None + def setUp(self): + self.client_mock = MagicMock(spec=weaviate.WeaviateClient) + self.client_mock.collections = MagicMock() + self.client_mock.collections.exists.return_value = False + self.client_mock.collections.create.return_value = None + self.client_mock.collections.delete.return_value = None -# patcher = patch('mem0.vector_stores.weaviate.weaviate.connect_to_local', return_value=self.client_mock) -# self.mock_weaviate = patcher.start() -# self.addCleanup(patcher.stop) + patcher = patch("mem0.vector_stores.weaviate.weaviate.connect_to_local", return_value=self.client_mock) + self.mock_weaviate = patcher.start() + self.addCleanup(patcher.stop) -# self.weaviate_db = Weaviate( -# collection_name="test_collection", -# embedding_model_dims=1536, -# cluster_url=os.getenv('WEAVIATE_CLUSTER_URL'), -# auth_client_secret=os.getenv('WEAVIATE_API_KEY'), -# additional_headers={"X-OpenAI-Api-Key": "test_key"}, -# ) + self.weaviate_db = Weaviate( + collection_name="test_collection", + embedding_model_dims=1536, + cluster_url=os.getenv("WEAVIATE_CLUSTER_URL"), + auth_client_secret=os.getenv("WEAVIATE_API_KEY"), + additional_headers={"X-OpenAI-Api-Key": "test_key"}, + ) -# self.client_mock.reset_mock() + self.client_mock.reset_mock() -# @classmethod -# def tearDownClass(cls): -# for key, value in cls.original_env.items(): -# if value is not None: -# os.environ[key] = value -# else: -# os.environ.pop(key, None) + @classmethod + def tearDownClass(cls): + for key, value in cls.original_env.items(): + if value is not None: + os.environ[key] = value + else: + os.environ.pop(key, None) -# def tearDown(self): -# self.client_mock.reset_mock() + def tearDown(self): + self.client_mock.reset_mock() -# def test_create_col(self): -# self.client_mock.collections.exists.return_value = False -# self.weaviate_db.create_col(vector_size=1536) + def test_create_col(self): + self.client_mock.collections.exists.return_value = False + self.weaviate_db.create_col(vector_size=1536) + + self.client_mock.collections.create.assert_called_once() + + self.client_mock.reset_mock() + + self.client_mock.collections.exists.return_value = True + self.weaviate_db.create_col(vector_size=1536) + + self.client_mock.collections.create.assert_not_called() + + def test_insert(self): + self.client_mock.batch = MagicMock() + + self.client_mock.batch.fixed_size.return_value.__enter__.return_value = MagicMock() + + self.client_mock.collections.get.return_value.data.insert_many.return_value = { + "results": [{"id": "id1"}, {"id": "id2"}] + } + + vectors = [[0.1] * 1536, [0.2] * 1536] + payloads = [{"key1": "value1"}, {"key2": "value2"}] + ids = [str(uuid.uuid4()), str(uuid.uuid4())] + + self.weaviate_db.insert(vectors=vectors, payloads=payloads, ids=ids) + + def test_get(self): + valid_uuid = str(uuid.uuid4()) + + mock_response = MagicMock() + mock_response.properties = { + "hash": "abc123", + "created_at": "2025-03-08T12:00:00Z", + "updated_at": "2025-03-08T13:00:00Z", + "user_id": "user_123", + "agent_id": "agent_456", + "run_id": "run_789", + "data": {"key": "value"}, + "category": "test", + } + mock_response.uuid = valid_uuid + + self.client_mock.collections.get.return_value.query.fetch_object_by_id.return_value = mock_response + + result = self.weaviate_db.get(vector_id=valid_uuid) + + assert result.id == valid_uuid + + expected_payload = mock_response.properties.copy() + expected_payload["id"] = valid_uuid + + assert result.payload == expected_payload + + def test_get_not_found(self): + mock_response = httpx.Response(status_code=404, json={"error": "Not found"}) + + self.client_mock.collections.get.return_value.data.get_by_id.side_effect = UnexpectedStatusCodeException( + "Not found", mock_response + ) + + def test_search(self): + mock_objects = [{"uuid": "id1", "properties": {"key1": "value1"}, "metadata": {"distance": 0.2}}] + + mock_response = MagicMock() + mock_response.objects = [] + + for obj in mock_objects: + mock_obj = MagicMock() + mock_obj.uuid = obj["uuid"] + mock_obj.properties = obj["properties"] + mock_obj.metadata = MagicMock() + mock_obj.metadata.distance = obj["metadata"]["distance"] + mock_response.objects.append(mock_obj) + + mock_hybrid = MagicMock() + self.client_mock.collections.get.return_value.query.hybrid = mock_hybrid + mock_hybrid.return_value = mock_response + + vectors = [[0.1] * 1536] + results = self.weaviate_db.search(query="", vectors=vectors, limit=5) + + mock_hybrid.assert_called_once() + + self.assertEqual(len(results), 1) + self.assertEqual(results[0].id, "id1") + self.assertEqual(results[0].score, 0.8) + + def test_delete(self): + self.weaviate_db.delete(vector_id="id1") + + self.client_mock.collections.get.return_value.data.delete_by_id.assert_called_once_with("id1") + + def test_list(self): + mock_objects = [] + + mock_obj1 = MagicMock() + mock_obj1.uuid = "id1" + mock_obj1.properties = {"key1": "value1"} + mock_objects.append(mock_obj1) + + mock_obj2 = MagicMock() + mock_obj2.uuid = "id2" + mock_obj2.properties = {"key2": "value2"} + mock_objects.append(mock_obj2) + + mock_response = MagicMock() + mock_response.objects = mock_objects + + mock_fetch = MagicMock() + self.client_mock.collections.get.return_value.query.fetch_objects = mock_fetch + mock_fetch.return_value = mock_response + + results = self.weaviate_db.list(limit=10) + + mock_fetch.assert_called_once() + + # Verify results + self.assertEqual(len(results), 1) + self.assertEqual(len(results[0]), 2) + self.assertEqual(results[0][0].id, "id1") + self.assertEqual(results[0][0].payload["key1"], "value1") + self.assertEqual(results[0][1].id, "id2") + self.assertEqual(results[0][1].payload["key2"], "value2") + + def test_list_cols(self): + mock_collection1 = MagicMock() + mock_collection1.name = "collection1" + + mock_collection2 = MagicMock() + mock_collection2.name = "collection2" + self.client_mock.collections.list_all.return_value = [mock_collection1, mock_collection2] + + result = self.weaviate_db.list_cols() + expected = {"collections": [{"name": "collection1"}, {"name": "collection2"}]} + + assert result == expected + + self.client_mock.collections.list_all.assert_called_once() + + def test_delete_col(self): + self.weaviate_db.delete_col() + + self.client_mock.collections.delete.assert_called_once_with("test_collection") -# self.client_mock.collections.create.assert_called_once() - - -# self.client_mock.reset_mock() - -# self.client_mock.collections.exists.return_value = True -# self.weaviate_db.create_col(vector_size=1536) - -# self.client_mock.collections.create.assert_not_called() - -# def test_insert(self): -# self.client_mock.batch = MagicMock() - -# self.client_mock.batch.fixed_size.return_value.__enter__.return_value = MagicMock() - -# self.client_mock.collections.get.return_value.data.insert_many.return_value = { -# "results": [{"id": "id1"}, {"id": "id2"}] -# } - -# vectors = [[0.1] * 1536, [0.2] * 1536] -# payloads = [{"key1": "value1"}, {"key2": "value2"}] -# ids = [str(uuid.uuid4()), str(uuid.uuid4())] - -# results = self.weaviate_db.insert(vectors=vectors, payloads=payloads, ids=ids) - -# def test_get(self): -# valid_uuid = str(uuid.uuid4()) - -# mock_response = MagicMock() -# mock_response.properties = { -# "hash": "abc123", -# "created_at": "2025-03-08T12:00:00Z", -# "updated_at": "2025-03-08T13:00:00Z", -# "user_id": "user_123", -# "agent_id": "agent_456", -# "run_id": "run_789", -# "data": {"key": "value"}, -# "category": "test", -# } -# mock_response.uuid = valid_uuid - -# self.client_mock.collections.get.return_value.query.fetch_object_by_id.return_value = mock_response - -# result = self.weaviate_db.get(vector_id=valid_uuid) - -# assert result.id == valid_uuid - -# expected_payload = mock_response.properties.copy() -# expected_payload["id"] = valid_uuid - -# assert result.payload == expected_payload - - -# def test_get_not_found(self): -# mock_response = httpx.Response(status_code=404, json={"error": "Not found"}) - -# self.client_mock.collections.get.return_value.data.get_by_id.side_effect = UnexpectedStatusCodeException( -# "Not found", mock_response -# ) - - -# def test_search(self): -# mock_objects = [ -# { -# "uuid": "id1", -# "properties": {"key1": "value1"}, -# "metadata": {"distance": 0.2} -# } -# ] - -# mock_response = MagicMock() -# mock_response.objects = [] - -# for obj in mock_objects: -# mock_obj = MagicMock() -# mock_obj.uuid = obj["uuid"] -# mock_obj.properties = obj["properties"] -# mock_obj.metadata = MagicMock() -# mock_obj.metadata.distance = obj["metadata"]["distance"] -# mock_response.objects.append(mock_obj) - -# mock_hybrid = MagicMock() -# self.client_mock.collections.get.return_value.query.hybrid = mock_hybrid -# mock_hybrid.return_value = mock_response - -# vectors = [[0.1] * 1536] -# results = self.weaviate_db.search(query="", vectors=vectors, limit=5) - -# mock_hybrid.assert_called_once() - -# self.assertEqual(len(results), 1) -# self.assertEqual(results[0].id, "id1") -# self.assertEqual(results[0].score, 0.8) - -# def test_delete(self): -# self.weaviate_db.delete(vector_id="id1") - -# self.client_mock.collections.get.return_value.data.delete_by_id.assert_called_once_with("id1") - -# def test_list(self): -# mock_objects = [] - -# mock_obj1 = MagicMock() -# mock_obj1.uuid = "id1" -# mock_obj1.properties = {"key1": "value1"} -# mock_objects.append(mock_obj1) - -# mock_obj2 = MagicMock() -# mock_obj2.uuid = "id2" -# mock_obj2.properties = {"key2": "value2"} -# mock_objects.append(mock_obj2) - -# mock_response = MagicMock() -# mock_response.objects = mock_objects - -# mock_fetch = MagicMock() -# self.client_mock.collections.get.return_value.query.fetch_objects = mock_fetch -# mock_fetch.return_value = mock_response - -# results = self.weaviate_db.list(limit=10) - -# mock_fetch.assert_called_once() - -# # Verify results -# self.assertEqual(len(results), 1) -# self.assertEqual(len(results[0]), 2) -# self.assertEqual(results[0][0].id, "id1") -# self.assertEqual(results[0][0].payload["key1"], "value1") -# self.assertEqual(results[0][1].id, "id2") -# self.assertEqual(results[0][1].payload["key2"], "value2") - - -# def test_list_cols(self): -# mock_collection1 = MagicMock() -# mock_collection1.name = "collection1" - -# mock_collection2 = MagicMock() -# mock_collection2.name = "collection2" -# self.client_mock.collections.list_all.return_value = [mock_collection1, mock_collection2] - -# result = self.weaviate_db.list_cols() -# expected = {"collections": [{"name": "collection1"}, {"name": "collection2"}]} - -# assert result == expected - -# self.client_mock.collections.list_all.assert_called_once() - - -# def test_delete_col(self): -# self.weaviate_db.delete_col() - -# self.client_mock.collections.delete.assert_called_once_with("test_collection") - - -# if __name__ == '__main__': -# unittest.main() +if __name__ == "__main__": + unittest.main()