From 73075c7d246b0e18caaaa0d251810c8e3862da88 Mon Sep 17 00:00:00 2001
From: Alexsander Hamir <alexsanderhamirgomesbaptista@gmail.com>
Date: Sat, 6 Dec 2025 08:11:52 -0800
Subject: [PATCH] fix: add retry logic for flaky Azure image generation health
 check test (#17595)

- Add missing @pytest.mark.asyncio decorator
- Implement retry logic with exponential backoff (3 retries)
- Only retry on transient Azure internal server errors
- Fail immediately on non-transient errors

This fixes the flaky test_azure_img_gen_health_check which was failing
due to transient Azure internal server errors that are outside our control.
---
 .../litellm_utils_tests/test_health_check.py  | 55 +++++++++++++++----
 1 file changed, 43 insertions(+), 12 deletions(-)

diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py
index 05c4598f82..98b1a35379 100644
--- a/tests/litellm_utils_tests/test_health_check.py
+++ b/tests/litellm_utils_tests/test_health_check.py
@@ -83,20 +83,51 @@ async def test_openai_img_gen_health_check():
 # asyncio.run(test_openai_img_gen_health_check())
 
 
+@pytest.mark.asyncio
 async def test_azure_img_gen_health_check():
+    """
+    Test Azure image generation health check with retry logic for transient errors.
+    Azure sometimes returns internal server errors which are transient and not something we can control.
+    """
     litellm._turn_on_debug()
-    response = await litellm.ahealth_check(
-        model_params={
-            "model": "azure/dall-e-3",
-            "api_base": os.getenv("AZURE_API_BASE"),
-            "api_key": os.getenv("AZURE_API_KEY"),
-        },
-        mode="image_generation",
-        prompt="cute baby sea otter",
-    )
-
-    assert isinstance(response, dict) and "error" not in response
-    return response
+    max_retries = 3
+    retry_delay = 1  # Start with 1 second delay
+    
+    for attempt in range(max_retries):
+        response = await litellm.ahealth_check(
+            model_params={
+                "model": "azure/dall-e-3",
+                "api_base": os.getenv("AZURE_API_BASE"),
+                "api_key": os.getenv("AZURE_API_KEY"),
+            },
+            mode="image_generation",
+            prompt="cute baby sea otter",
+        )
+        
+        # Check if response is successful (no error)
+        if isinstance(response, dict) and "error" not in response:
+            return response
+        
+        # Check if error is a transient Azure internal server error
+        error_str = str(response.get("error", "")).lower()
+        is_transient_error = (
+            "internalservererror" in error_str
+            or "internal server error" in error_str
+            or "internalfailure" in error_str
+            or "internal failure" in error_str
+        )
+        
+        # If it's the last attempt or not a transient error, fail the test
+        if attempt == max_retries - 1 or not is_transient_error:
+            assert isinstance(response, dict) and "error" not in response, f"Health check failed: {response.get('error', 'Unknown error')}"
+            return response
+        
+        # Wait before retrying with exponential backoff
+        await asyncio.sleep(retry_delay)
+        retry_delay *= 2  # Exponential backoff
+    
+    # Should not reach here, but just in case
+    assert False, "Health check failed after all retries"
 
 
 @pytest.mark.skip(reason="AWS Suspended Account")