Spaces:
Sleeping
Sleeping
Yannick Lemin
junie made this beautiful block of code, it's not working, that's tomorrow's problem
1fed057
| import unittest | |
| from unittest.mock import patch, MagicMock | |
| from src.search.vector_search import VectorSearch | |
| from src.utils.constants import Constants | |
| class TestVectorSearch(unittest.TestCase): | |
| """ | |
| Test cases for the VectorSearch class. | |
| """ | |
| def test_init_success(self, mock_db_connection, mock_vector_search, mock_embeddings, | |
| mock_getenv, mock_load_dotenv): | |
| """ | |
| Test successful initialization of VectorSearch. | |
| """ | |
| # Setup mocks | |
| mock_getenv.side_effect = lambda key: "test_key" if key == "OPENAI_API_KEY" else "mongodb://test_uri" | |
| mock_db_instance = MagicMock() | |
| mock_collection = MagicMock() | |
| mock_db_instance.get_collection.return_value = mock_collection | |
| mock_db_connection.return_value = mock_db_instance | |
| mock_embeddings_instance = MagicMock() | |
| mock_embeddings.return_value = mock_embeddings_instance | |
| mock_vector_search_instance = MagicMock() | |
| mock_vector_search.from_connection_string.return_value = mock_vector_search_instance | |
| # Create instance | |
| vector_search = VectorSearch() | |
| # Assertions | |
| mock_load_dotenv.assert_called_once() | |
| mock_getenv.assert_any_call("OPENAI_API_KEY") | |
| mock_db_connection.assert_called_once() | |
| mock_db_instance.get_collection.assert_called_once() | |
| mock_embeddings.assert_called_once_with( | |
| model=Constants.EMBEDDING_MODEL, | |
| openai_api_key="test_key" | |
| ) | |
| mock_vector_search.from_connection_string.assert_called_once_with( | |
| connection_string="mongodb://test_uri", | |
| namespace=f"{Constants.DB_NAME}.{Constants.COLLECTION_NAME}", | |
| embedding=mock_embeddings_instance, | |
| index_name=Constants.ATLAS_VECTOR_SEARCH_INDEX_NAME, | |
| text_key="text" | |
| ) | |
| self.assertEqual(vector_search.vector_search, mock_vector_search_instance) | |
| def test_init_missing_api_key(self, mock_getenv, mock_load_dotenv): | |
| """ | |
| Test initialization with missing OpenAI API key. | |
| """ | |
| # Setup mocks | |
| mock_getenv.return_value = None | |
| # Assert that ValueError is raised | |
| with self.assertRaises(ValueError) as context: | |
| VectorSearch() | |
| self.assertEqual(str(context.exception), "OPENAI_API_KEY environment variable is not set") | |
| def test_search(self, mock_db_connection, mock_vector_search, mock_embeddings, | |
| mock_getenv, mock_load_dotenv): | |
| """ | |
| Test search method. | |
| """ | |
| # Setup mocks | |
| mock_getenv.side_effect = lambda key: "test_key" if key == "OPENAI_API_KEY" else "mongodb://test_uri" | |
| mock_db_instance = MagicMock() | |
| mock_db_connection.return_value = mock_db_instance | |
| mock_vector_search_instance = MagicMock() | |
| mock_vector_search.from_connection_string.return_value = mock_vector_search_instance | |
| # Create mock search results | |
| mock_doc1 = MagicMock() | |
| mock_doc1.page_content = "Test content 1" | |
| mock_doc1.metadata = { | |
| "id": "1", | |
| "title": "Test Title 1", | |
| "price": "100€", | |
| "thumbnail": "http://example.com/image1.jpg", | |
| "product_page_url": "http://example.com/product1" | |
| } | |
| mock_doc2 = MagicMock() | |
| mock_doc2.page_content = "Test content 2" | |
| mock_doc2.metadata = { | |
| "id": "2", | |
| "title": "Test Title 2", | |
| "price": "200€", | |
| "thumbnail": "", | |
| "product_page_url": "http://example.com/product2" | |
| } | |
| mock_vector_search_instance.similarity_search.return_value = [mock_doc1, mock_doc2] | |
| # Create instance and call method | |
| vector_search = VectorSearch() | |
| results = vector_search.search("test query", limit=2) | |
| # Assertions | |
| mock_vector_search_instance.similarity_search.assert_called_once_with("test query", k=2) | |
| self.assertEqual(len(results), 2) | |
| self.assertEqual(results[0]["id"], "1") | |
| self.assertEqual(results[0]["title"], "Test Title 1") | |
| self.assertEqual(results[0]["text"], "Test content 1") | |
| self.assertEqual(results[0]["price"], "100€") | |
| self.assertEqual(results[0]["thumbnail"], "http://example.com/image1.jpg") | |
| self.assertEqual(results[0]["product_page_url"], "http://example.com/product1") | |
| self.assertEqual(results[1]["id"], "2") | |
| self.assertEqual(results[1]["title"], "Test Title 2") | |
| self.assertEqual(results[1]["text"], "Test content 2") | |
| self.assertEqual(results[1]["price"], "200€") | |
| self.assertEqual(results[1]["thumbnail"], Constants.PLACEHOLDER_IMAGE_URL) | |
| self.assertEqual(results[1]["product_page_url"], "http://example.com/product2") | |
| def test_close(self, mock_db_connection, mock_vector_search, mock_embeddings, | |
| mock_getenv, mock_load_dotenv): | |
| """ | |
| Test close method. | |
| """ | |
| # Setup mocks | |
| mock_getenv.side_effect = lambda key: "test_key" if key == "OPENAI_API_KEY" else "mongodb://test_uri" | |
| mock_db_instance = MagicMock() | |
| mock_db_connection.return_value = mock_db_instance | |
| # Create instance and call method | |
| vector_search = VectorSearch() | |
| vector_search.close() | |
| # Assertions | |
| mock_db_instance.close_connection.assert_called_once() | |
| if __name__ == '__main__': | |
| unittest.main() |