From 74ef034110bc56bae5a8260b6f14ae4091f69df3 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 21 Feb 2026 11:56:29 -0800
Subject: [PATCH] fix(tests): add flaky retries to flaky CI tests (#21795)

* fix(tests): add flaky retries and error handling to test_create_eval

* fix(tests): add flaky retries to test_cohere_v2_conversation_history

* fix(tests): add flaky retries to test_gemini_url_context
---
 tests/llm_translation/test_cohere.py    | 15 ++++---
 tests/llm_translation/test_evals_api.py | 58 ++++++++++++++-----------
 tests/llm_translation/test_gemini.py    |  1 +
 3 files changed, 41 insertions(+), 33 deletions(-)

diff --git a/tests/llm_translation/test_cohere.py b/tests/llm_translation/test_cohere.py
index 2d7e5da20d..6f6266c6a0 100644
--- a/tests/llm_translation/test_cohere.py
+++ b/tests/llm_translation/test_cohere.py
@@ -837,6 +837,7 @@ async def test_cohere_documents_options_in_request_body():
 
 
 @pytest.mark.asyncio
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_cohere_v2_conversation_history():
     """Test Cohere v2 with conversation history."""
     try:
@@ -847,20 +848,20 @@ async def test_cohere_v2_conversation_history():
             {"role": "assistant", "content": "2+2 equals 4."},
             {"role": "user", "content": "What about 3+3?"}
         ]
-        
+
         response = await litellm.acompletion(
             model="cohere_chat/v2/command-a-03-2025",
             messages=messages,
             max_tokens=50
         )
-        
+
         # Validate response with conversation history
         assert response.choices is not None
         assert len(response.choices) > 0
         assert response.choices[0].message.content is not None
         print(f"Conversation history response: {response.choices[0].message.content}")
-        
-    except litellm.ServiceUnavailableError:
-        pass
-    except Exception as e:
-        pytest.fail(f"Error occurred: {e}")
\ No newline at end of file
+
+    except (litellm.ServiceUnavailableError, litellm.InternalServerError, litellm.Timeout, litellm.APIConnectionError):
+        pytest.skip("Cohere service unavailable")
+    except litellm.RateLimitError:
+        pytest.skip("Rate limit exceeded")
\ No newline at end of file
diff --git a/tests/llm_translation/test_evals_api.py b/tests/llm_translation/test_evals_api.py
index cc8f7897cd..945186249f 100644
--- a/tests/llm_translation/test_evals_api.py
+++ b/tests/llm_translation/test_evals_api.py
@@ -41,6 +41,7 @@ class BaseEvalsAPITest(ABC):
         """Return the API base URL for the provider"""
         pass
 
+    @pytest.mark.flaky(retries=3, delay=2)
     def test_create_eval(self):
         """
         Test creating an evaluation.
@@ -59,32 +60,37 @@ class BaseEvalsAPITest(ABC):
         # Create eval with stored_completions data source
         unique_name = f"Test Eval {int(time.time())}"
 
-        response = litellm.create_eval(
-            name=unique_name,
-            data_source_config={
-                "type": "stored_completions",
-                "metadata": {"usecase": "chatbot"},
-            },
-            testing_criteria=[
-                {
-                    "type": "label_model",
-                    "model": "gpt-4o",
-                    "input": [
-                        {
-                            "role": "developer",
-                            "content": "Classify the sentiment as 'positive' or 'negative'",
-                        },
-                        {"role": "user", "content": "Statement: {{item.input}}"},
-                    ],
-                    "passing_labels": ["positive"],
-                    "labels": ["positive", "negative"],
-                    "name": "Sentiment grader",
-                }
-            ],
-            custom_llm_provider=custom_llm_provider,
-            api_key=api_key,
-            api_base=api_base,
-        )
+        try:
+            response = litellm.create_eval(
+                name=unique_name,
+                data_source_config={
+                    "type": "stored_completions",
+                    "metadata": {"usecase": "chatbot"},
+                },
+                testing_criteria=[
+                    {
+                        "type": "label_model",
+                        "model": "gpt-4o",
+                        "input": [
+                            {
+                                "role": "developer",
+                                "content": "Classify the sentiment as 'positive' or 'negative'",
+                            },
+                            {"role": "user", "content": "Statement: {{item.input}}"},
+                        ],
+                        "passing_labels": ["positive"],
+                        "labels": ["positive", "negative"],
+                        "name": "Sentiment grader",
+                    }
+                ],
+                custom_llm_provider=custom_llm_provider,
+                api_key=api_key,
+                api_base=api_base,
+            )
+        except (litellm.InternalServerError, litellm.APIConnectionError, litellm.Timeout, litellm.ServiceUnavailableError):
+            pytest.skip("Provider service unavailable")
+        except litellm.RateLimitError:
+            pytest.skip("Rate limit exceeded")
 
         assert response is not None
         assert isinstance(response, Eval)
diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py
index c1c52757cf..3c8ea49ba8 100644
--- a/tests/llm_translation/test_gemini.py
+++ b/tests/llm_translation/test_gemini.py
@@ -489,6 +489,7 @@ def test_gemini_finish_reason():
     assert response.choices[0].finish_reason == "length"
 
 
+@pytest.mark.flaky(retries=3, delay=2)
 def test_gemini_url_context():
     from litellm import completion