fix flaky tests

tests/test_litellm/containers/test_container_api.py

@@ -179,8 +179,8 @@ class TestContainerAPI:
             name="Retrieved Container"
         )
         
-        with patch('litellm.containers.main.base_llm_http_handler') as mock_handler:
-            mock_handler.container_retrieve_handler.return_value = mock_response
+        with patch('litellm.containers.main.base_llm_http_handler.container_retrieve_handler') as mock_handler:
+            mock_handler.return_value = mock_response
             
             response = retrieve_container(
                 container_id=container_id,

tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_resend_email.py

@@ -72,28 +72,38 @@ async def test_send_email_success(mock_env_vars, mock_httpx_client):
 
 @pytest.mark.asyncio
 async def test_send_email_missing_api_key(mock_httpx_client):
-    # Remove the API key from environment
-    if "RESEND_API_KEY" in os.environ:
-        del os.environ["RESEND_API_KEY"]
+    # Remove the API key from environment before initializing logger
+    original_key = os.environ.pop("RESEND_API_KEY", None)
+    
+    try:
+        # Initialize the logger after removing the API key
+        logger = ResendEmailLogger()
 
-    # Initialize the logger
-    logger = ResendEmailLogger()
+        # Test data
+        from_email = "test@example.com"
+        to_email = ["recipient@example.com"]
+        subject = "Test Subject"
+        html_body = "<p>Test email body</p>"
 
-    # Test data
-    from_email = "test@example.com"
-    to_email = ["recipient@example.com"]
-    subject = "Test Subject"
-    html_body = "<p>Test email body</p>"
+        # Mock the response to avoid making real HTTP requests
+        mock_response = mock.AsyncMock(spec=Response)
+        mock_response.status_code = 200
+        mock_response.json.return_value = {"id": "test_email_id"}
+        mock_httpx_client.post.return_value = mock_response
 
-    # Send email
-    await logger.send_email(
-        from_email=from_email, to_email=to_email, subject=subject, html_body=html_body
-    )
+        # Send email
+        await logger.send_email(
+            from_email=from_email, to_email=to_email, subject=subject, html_body=html_body
+        )
 
-    # Verify the HTTP client was called with None as the API key
-    mock_httpx_client.post.assert_called_once()
-    call_args = mock_httpx_client.post.call_args
-    assert call_args[1]["headers"] == {"Authorization": "Bearer None"}
+        # Verify the HTTP client was called with None as the API key
+        mock_httpx_client.post.assert_called_once()
+        call_args = mock_httpx_client.post.call_args
+        assert call_args[1]["headers"] == {"Authorization": "Bearer None"}
+    finally:
+        # Restore the original key if it existed
+        if original_key is not None:
+            os.environ["RESEND_API_KEY"] = original_key
 
 
 @pytest.mark.asyncio
@@ -107,6 +117,12 @@ async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client):
     subject = "Test Subject"
     html_body = "<p>Test email body</p>"
 
+    # Mock the response to avoid making real HTTP requests
+    mock_response = mock.AsyncMock(spec=Response)
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"id": "test_email_id"}
+    mock_httpx_client.post.return_value = mock_response
+
     # Send email
     await logger.send_email(
         from_email=from_email, to_email=to_email, subject=subject, html_body=html_body

tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_sendgrid_email.py

@@ -63,9 +63,17 @@ async def test_send_email_success(mock_env_vars, mock_httpx_client):
 
 @pytest.mark.asyncio
 async def test_send_email_missing_api_key(mock_httpx_client):
-    with mock.patch.dict(os.environ, {}, clear=True):
+    # Remove the API key from environment before initializing logger
+    original_key = os.environ.pop("SENDGRID_API_KEY", None)
+    
+    try:
         logger = SendGridEmailLogger()
 
+        # Mock the response to avoid making real HTTP requests
+        mock_response = mock.AsyncMock(spec=Response)
+        mock_response.status_code = 401
+        mock_httpx_client.post.return_value = mock_response
+
         with pytest.raises(ValueError):
             await logger.send_email(
                 from_email="test@example.com",
@@ -75,6 +83,10 @@ async def test_send_email_missing_api_key(mock_httpx_client):
             )
 
         mock_httpx_client.post.assert_not_called()
+    finally:
+        # Restore the original key if it existed
+        if original_key is not None:
+            os.environ["SENDGRID_API_KEY"] = original_key
 
 
 @pytest.mark.asyncio
@@ -86,6 +98,12 @@ async def test_send_email_multiple_recipients(mock_env_vars, mock_httpx_client):
     subject = "Test Subject"
     html_body = "<p>Test email body</p>"
 
+    # Mock the response to avoid making real HTTP requests
+    mock_response = mock.AsyncMock(spec=Response)
+    mock_response.status_code = 202
+    mock_response.text = "accepted"
+    mock_httpx_client.post.return_value = mock_response
+
     await logger.send_email(
         from_email=from_email, to_email=to_email, subject=subject, html_body=html_body
     )

tests/test_litellm/test_main.py

@@ -488,8 +488,11 @@ async def test_openai_env_base(
     model = "gpt-4o"
     messages = [{"role": "user", "content": "Hello, how are you?"}]
 
-    # Ensure respx_mock is properly configured
-    respx_mock.route(host="localhost", port=12345).post("/v1/chat/completions").respond(
+    # Ensure respx_mock is properly configured - use correct respx API
+    respx_mock.post(
+        url__regex=r"http://localhost:12345/v1/chat/completions.*"
+    ).mock(return_value=httpx.Response(
+        status_code=200,
         json={
             "id": "chatcmpl-123",
             "object": "chat.completion",
@@ -507,7 +510,7 @@ async def test_openai_env_base(
             ],
             "usage": {"prompt_tokens": 9, "completion_tokens": 12, "total_tokens": 21},
         }
-    )
+    ))
 
     response = await litellm.acompletion(model=model, messages=messages)
 

tests/test_litellm/vector_stores/test_vector_store_registry.py

@@ -133,6 +133,16 @@ def test_search_uses_registry_credentials():
     try:
         logger = MagicMock()
         logger._response_cost_calculator.return_value = 0
+        
+        # Mock the search response
+        mock_search_response = {
+            "object": "list",
+            "data": [],
+            "first_id": None,
+            "last_id": None,
+            "has_more": False
+        }
+        
         with patch.object(
             registry,
             "get_credentials_for_vector_store",
@@ -142,7 +152,7 @@ def test_search_uses_registry_credentials():
             return_value=MagicMock(),
         ), patch(
             "litellm.vector_stores.main.base_llm_http_handler.vector_store_search_handler",
-            return_value={},
+            return_value=mock_search_response,
         ) as mock_handler:
             search(vector_store_id="vs1", query="test", litellm_logging_obj=logger)
             mock_get_creds.assert_called_once_with("vs1")