Spaces:
Sleeping
Sleeping
File size: 6,609 Bytes
1fed057 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
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.
"""
@patch('src.search.vector_search.load_dotenv')
@patch('src.search.vector_search.os.getenv')
@patch('src.search.vector_search.OpenAIEmbeddings')
@patch('src.search.vector_search.MongoDBAtlasVectorSearch')
@patch('src.search.vector_search.DatabaseConnection')
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)
@patch('src.search.vector_search.load_dotenv')
@patch('src.search.vector_search.os.getenv')
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")
@patch('src.search.vector_search.load_dotenv')
@patch('src.search.vector_search.os.getenv')
@patch('src.search.vector_search.OpenAIEmbeddings')
@patch('src.search.vector_search.MongoDBAtlasVectorSearch')
@patch('src.search.vector_search.DatabaseConnection')
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")
@patch('src.search.vector_search.load_dotenv')
@patch('src.search.vector_search.os.getenv')
@patch('src.search.vector_search.OpenAIEmbeddings')
@patch('src.search.vector_search.MongoDBAtlasVectorSearch')
@patch('src.search.vector_search.DatabaseConnection')
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() |