44
55import pytest
66import src .agentic_rag .tools as tools_module
7- from dotenv import load_dotenv
87
98# Add parent directory to path
109sys .path .insert (0 , os .path .abspath (os .path .join (os .path .dirname (__file__ ), ".." )))
@@ -177,19 +176,16 @@ def test_get_retriever_components_initialization(mock_get_env, mock_client_class
177176 tools_module ._client_cache = None
178177 tools_module ._vector_store_id_cache = None
179178
180- # Mock environment variable
181- mock_get_env .return_value = "http://localhost:8321"
179+ # Mock environment variables: BASE_URL, VECTOR_STORE_ID, API_KEY
180+ def getenv_side_effect (key ):
181+ return {
182+ "BASE_URL" : "http://localhost:8321" ,
183+ "VECTOR_STORE_ID" : "test-vector-store-123" ,
184+ "API_KEY" : "test-key" ,
185+ }.get (key )
182186
183- # Mock client and vector store list
184- mock_client = Mock ()
185- mock_vector_store = Mock ()
186- mock_vector_store .id = "test-vector-store-123"
187-
188- mock_list_response = Mock ()
189- mock_list_response .data = [mock_vector_store ]
190-
191- mock_client .vector_stores .list .return_value = mock_list_response
192- mock_client_class .return_value = mock_client
187+ mock_get_env .side_effect = getenv_side_effect
188+ mock_client_class .return_value = Mock ()
193189
194190 # Call function
195191 result = get_retriever_components ()
@@ -198,7 +194,9 @@ def test_get_retriever_components_initialization(mock_get_env, mock_client_class
198194 assert "client" in result
199195 assert "vector_store_id" in result
200196 assert result ["vector_store_id" ] == "test-vector-store-123"
201- mock_client_class .assert_called_once_with (base_url = "http://localhost:8321" )
197+ mock_client_class .assert_called_once_with (
198+ base_url = "http://localhost:8321" , api_key = "test-key"
199+ )
202200
203201
204202@patch ("src.agentic_rag.tools.LlamaStackClient" )
@@ -220,62 +218,51 @@ def test_get_retriever_components_caching(mock_get_env, mock_client_class):
220218
221219
222220@patch ("src.agentic_rag.tools.LlamaStackClient" )
223- def test_get_retriever_components_with_base_url (mock_client_class ):
221+ @patch ("src.agentic_rag.tools.getenv" )
222+ def test_get_retriever_components_with_base_url (mock_get_env , mock_client_class ):
224223 """Test that base_url parameter is used when provided."""
225224 # Reset cache
226225 tools_module ._client_cache = None
227226 tools_module ._vector_store_id_cache = None
228227
229- # Mock client and vector store list
230- mock_client = Mock ()
231- mock_vector_store = Mock ()
232- mock_vector_store .id = "test-id"
233-
234- mock_list_response = Mock ()
235- mock_list_response .data = [mock_vector_store ]
228+ def getenv_side_effect (key ):
229+ return {
230+ "VECTOR_STORE_ID" : "test-id" ,
231+ "API_KEY" : "test-key" ,
232+ }.get (key )
236233
237- mock_client . vector_stores . list . return_value = mock_list_response
238- mock_client_class .return_value = mock_client
234+ mock_get_env . side_effect = getenv_side_effect
235+ mock_client_class .return_value = Mock ()
239236
240237 # Call with explicit base_url
241238 result = get_retriever_components (base_url = "http://custom:9999" )
242239
243- # Should use provided base_url
244- mock_client_class .assert_called_once_with (base_url = "http://custom:9999" )
240+ # Should use provided base_url (stripped of /v1 suffix if present)
241+ mock_client_class .assert_called_once_with (
242+ base_url = "http://custom:9999" , api_key = "test-key"
243+ )
245244 assert result ["vector_store_id" ] == "test-id"
246245
247246
248- @patch ("src.agentic_rag.tools.LlamaStackClient" )
249247@patch ("src.agentic_rag.tools.getenv" )
250- def test_get_retriever_components_no_vector_store (mock_get_env , mock_client_class ):
251- """Test error handling when no vector store is found ."""
248+ def test_get_retriever_components_no_vector_store (mock_get_env ):
249+ """Test error handling when VECTOR_STORE_ID env var is not set ."""
252250 # Reset cache
253251 tools_module ._client_cache = None
254252 tools_module ._vector_store_id_cache = None
255253
256- mock_get_env .return_value = "http://localhost:8321"
257-
258- # Mock client with empty vector store list
259- mock_client = Mock ()
260- mock_list_response = Mock ()
261- mock_list_response .data = [] # No vector stores
254+ def getenv_side_effect (key ):
255+ return {"BASE_URL" : "http://localhost:8321" }.get (key )
262256
263- mock_client .vector_stores .list .return_value = mock_list_response
264- mock_client_class .return_value = mock_client
257+ mock_get_env .side_effect = getenv_side_effect
265258
266- # Should raise RuntimeError
259+ # Should raise RuntimeError when VECTOR_STORE_ID is missing
267260 with pytest .raises (RuntimeError ) as exc_info :
268261 get_retriever_components ()
269262
270- assert "No vector store found " in str (exc_info .value )
263+ assert "VECTOR_STORE_ID " in str (exc_info .value )
271264 assert "load_documents.py" in str (exc_info .value )
272265
273266
274- def test_get_retriever_components ():
275- load_dotenv (verbose = True )
276- base_url = os .getenv ("BASE_URL" )
277- get_retriever_components (base_url )
278-
279-
280267if __name__ == "__main__" :
281268 pytest .main ([__file__ , "-v" ])
0 commit comments