From 3966c545b4d3acb510f56c3507c9777a60c074f6 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 27 Aug 2024 22:47:11 -0700
Subject: [PATCH] fix flaky tests

---
 litellm/tests/test_completion.py               | 1 +
 litellm/tests/test_parallel_request_limiter.py | 1 +
 litellm/tests/test_router_caching.py           | 1 +
 litellm/tests/test_streaming.py                | 1 +
 4 files changed, 4 insertions(+)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 01435d6a3e0..baf528d8bc5 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -4078,6 +4078,7 @@ def test_completion_nvidia_nim():
         # "gemini-1.5-flash",
     ],
 )
+@pytest.mark.flaky(retries=3, delay=1)
 def test_completion_gemini(model):
     litellm.set_verbose = True
     model_name = "gemini/{}".format(model)
diff --git a/litellm/tests/test_parallel_request_limiter.py b/litellm/tests/test_parallel_request_limiter.py
index 463f9cf50d1..5173bc9c058 100644
--- a/litellm/tests/test_parallel_request_limiter.py
+++ b/litellm/tests/test_parallel_request_limiter.py
@@ -758,6 +758,7 @@ async def test_streaming_router_tpm_limit():
 
 
 @pytest.mark.asyncio
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_bad_router_call():
     litellm.set_verbose = True
     model_list = [
diff --git a/litellm/tests/test_router_caching.py b/litellm/tests/test_router_caching.py
index cfe3bc7c0b1..10b23c295bf 100644
--- a/litellm/tests/test_router_caching.py
+++ b/litellm/tests/test_router_caching.py
@@ -142,6 +142,7 @@ async def test_acompletion_caching_on_router():
 
 
 @pytest.mark.asyncio
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_completion_caching_on_router():
     # tests completion + caching on router
     try:
diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index d7b6af8e1aa..d2ef8aafc7b 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -714,6 +714,7 @@ async def test_acompletion_claude_2_stream():
     [True, False],
 )  # ,
 @pytest.mark.asyncio
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_completion_gemini_stream(sync_mode):
     try:
         litellm.set_verbose = True