my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	# ruff: noqa: E501
#2
#3	import os
#4	from unittest import mock
#5	from unittest.mock import Mock, patch
#6
#7	import pytest
#8
#9	from embedchain.config import ZillizDBConfig
#10	from embedchain.vectordb.zilliz import ZillizVectorDB
#11
#12
#13	# to run tests, provide the URI and TOKEN in .env file
#14	class TestZillizVectorDBConfig:
#15	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#16	def test_init_with_uri_and_token(self):
#17	"""
#18	Test if the `ZillizVectorDBConfig` instance is initialized with the correct uri and token values.
#19	"""
#20	# Create a ZillizDBConfig instance with mocked values
#21	expected_uri = "mocked_uri"
#22	expected_token = "mocked_token"
#23	db_config = ZillizDBConfig()
#24
#25	# Assert that the values in the ZillizVectorDB instance match the mocked values
#26	assert db_config.uri == expected_uri
#27	assert db_config.token == expected_token
#28
#29	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#30	def test_init_without_uri(self):
#31	"""
#32	Test if the `ZillizVectorDBConfig` instance throws an error when no URI found.
#33	"""
#34	try:
#35	del os.environ["ZILLIZ_CLOUD_URI"]
#36	except KeyError:
#37	pass
#38
#39	with pytest.raises(AttributeError):
#40	ZillizDBConfig()
#41
#42	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#43	def test_init_without_token(self):
#44	"""
#45	Test if the `ZillizVectorDBConfig` instance throws an error when no Token found.
#46	"""
#47	try:
#48	del os.environ["ZILLIZ_CLOUD_TOKEN"]
#49	except KeyError:
#50	pass
#51	# Test if an exception is raised when ZILLIZ_CLOUD_TOKEN is missing
#52	with pytest.raises(AttributeError):
#53	ZillizDBConfig()
#54
#55
#56	class TestZillizVectorDB:
#57	@pytest.fixture
#58	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#59	def mock_config(self, mocker):
#60	return mocker.Mock(spec=ZillizDBConfig())
#61
#62	@patch("embedchain.vectordb.zilliz.MilvusClient", autospec=True)
#63	@patch("embedchain.vectordb.zilliz.connections.connect", autospec=True)
#64	def test_zilliz_vector_db_setup(self, mock_connect, mock_client, mock_config):
#65	"""
#66	Test if the `ZillizVectorDB` instance is initialized with the correct uri and token values.
#67	"""
#68	# Create an instance of ZillizVectorDB with the mock config
#69	# zilliz_db = ZillizVectorDB(config=mock_config)
#70	ZillizVectorDB(config=mock_config)
#71
#72	# Assert that the MilvusClient and connections.connect were called
#73	mock_client.assert_called_once_with(uri=mock_config.uri, token=mock_config.token)
#74	mock_connect.assert_called_once_with(uri=mock_config.uri, token=mock_config.token)
#75
#76
#77	class TestZillizDBCollection:
#78	@pytest.fixture
#79	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#80	def mock_config(self, mocker):
#81	return mocker.Mock(spec=ZillizDBConfig())
#82
#83	@pytest.fixture
#84	def mock_embedder(self, mocker):
#85	return mocker.Mock()
#86
#87	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#88	def test_init_with_default_collection(self):
#89	"""
#90	Test if the `ZillizVectorDB` instance is initialized with the correct default collection name.
#91	"""
#92	# Create a ZillizDBConfig instance
#93	db_config = ZillizDBConfig()
#94
#95	assert db_config.collection_name == "embedchain_store"
#96
#97	@mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
#98	def test_init_with_custom_collection(self):
#99	"""
#100	Test if the `ZillizVectorDB` instance is initialized with the correct custom collection name.
#101	"""
#102	# Create a ZillizDBConfig instance with mocked values
#103
#104	expected_collection = "test_collection"
#105	db_config = ZillizDBConfig(collection_name="test_collection")
#106
#107	assert db_config.collection_name == expected_collection
#108
#109	@patch("embedchain.vectordb.zilliz.MilvusClient", autospec=True)
#110	@patch("embedchain.vectordb.zilliz.connections", autospec=True)
#111	def test_query(self, mock_connect, mock_client, mock_embedder, mock_config):
#112	# Create an instance of ZillizVectorDB with mock config
#113	zilliz_db = ZillizVectorDB(config=mock_config)
#114
#115	# Add a 'embedder' attribute to the ZillizVectorDB instance for testing
#116	zilliz_db.embedder = mock_embedder # Mock the 'collection' object
#117
#118	# Add a 'collection' attribute to the ZillizVectorDB instance for testing
#119	zilliz_db.collection = Mock(is_empty=False) # Mock the 'collection' object
#120
#121	assert zilliz_db.client == mock_client()
#122
#123	# Mock the MilvusClient search method
#124	with patch.object(zilliz_db.client, "search") as mock_search:
#125	# Mock the embedding function
#126	mock_embedder.embedding_fn.return_value = ["query_vector"]
#127
#128	# Mock the search result
#129	mock_search.return_value = [
#130	[
#131	{
#132	"distance": 0.0,
#133	"entity": {
#134	"text": "result_doc",
#135	"embeddings": [1, 2, 3],
#136	"metadata": {"url": "url_1", "doc_id": "doc_id_1"},
#137	},
#138	}
#139	]
#140	]
#141
#142	query_result = zilliz_db.query(input_query="query_text", n_results=1, where={})
#143
#144	# Assert that MilvusClient.search was called with the correct parameters
#145	mock_search.assert_called_with(
#146	collection_name=mock_config.collection_name,
#147	data=["query_vector"],
#148	filter="",
#149	limit=1,
#150	output_fields=["*"],
#151	)
#152
#153	# Assert that the query result matches the expected result
#154	assert query_result == ["result_doc"]
#155
#156	query_result_with_citations = zilliz_db.query(
#157	input_query="query_text", n_results=1, where={}, citations=True
#158	)
#159
#160	mock_search.assert_called_with(
#161	collection_name=mock_config.collection_name,
#162	data=["query_vector"],
#163	filter="",
#164	limit=1,
#165	output_fields=["*"],
#166	)
#167
#168	assert query_result_with_citations == [("result_doc", {"url": "url_1", "doc_id": "doc_id_1", "score": 0.0})]
#169

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public