From 32e445c59d46c3c8c336dc83aca4d6dce6d61bbd Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 11 May 2024 19:10:33 -0700
Subject: [PATCH] fix - unit tests for router retries

---
 litellm/tests/test_router_retries.py | 54 +++-------------------------
 1 file changed, 4 insertions(+), 50 deletions(-)

diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py
index f8109043ef..3a89b644b5 100644
--- a/litellm/tests/test_router_retries.py
+++ b/litellm/tests/test_router_retries.py
@@ -279,7 +279,6 @@ def test_retry_rate_limit_error_with_healthy_deployments():
         "deployment1",
         "deployment2",
     ]  # multiple healthy deployments mocked up
-    fallbacks = None
 
     router = litellm.Router(
         model_list=[
@@ -298,7 +297,7 @@ def test_retry_rate_limit_error_with_healthy_deployments():
     # Act & Assert
     try:
         response = router.should_retry_this_error(
-            rate_limit_error, healthy_deployments, fallbacks
+            error=rate_limit_error, healthy_deployments=healthy_deployments
         )
         print("response from should_retry_this_error: ", response)
     except Exception as e:
@@ -313,7 +312,6 @@ def test_do_not_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployme
     Test 2. It SHOULD NOT Retry, when healthy_deployments is [] and fallbacks is None
     """
     healthy_deployments = []
-    fallbacks = None
 
     router = litellm.Router(
         model_list=[
@@ -332,7 +330,7 @@ def test_do_not_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployme
     # Act & Assert
     try:
         response = router.should_retry_this_error(
-            rate_limit_error, healthy_deployments, fallbacks
+            error=rate_limit_error, healthy_deployments=healthy_deployments
         )
         assert response != True, "Should have raised RateLimitError"
     except openai.RateLimitError:
@@ -352,7 +350,7 @@ def test_raise_context_window_exceeded_error():
         llm_provider="azure",
         model="gpt-3.5-turbo",
     )
-    context_window_fallbacks = ["fallback1", "fallback2"]
+    context_window_fallbacks = [{"gpt-3.5-turbo": ["azure/chatgpt-v-2"]}]
 
     router = litellm.Router(
         model_list=[
@@ -371,7 +369,6 @@ def test_raise_context_window_exceeded_error():
     response = router.should_retry_this_error(
         error=context_window_error,
         healthy_deployments=None,
-        fallbacks=None,
         context_window_fallbacks=context_window_fallbacks,
     )
     assert (
@@ -412,7 +409,6 @@ def test_raise_context_window_exceeded_error_no_retry():
         response = router.should_retry_this_error(
             error=context_window_error,
             healthy_deployments=None,
-            fallbacks=None,
             context_window_fallbacks=context_window_fallbacks,
         )
         assert (
@@ -460,7 +456,6 @@ def test_timeout_for_rate_limit_error_with_healthy_deployments():
         remaining_retries=4,
         num_retries=4,
         healthy_deployments=healthy_deployments,
-        fallbacks=fallbacks,
     )
 
     print(
@@ -473,51 +468,11 @@ def test_timeout_for_rate_limit_error_with_healthy_deployments():
     assert _timeout == 0.0
 
 
-def test_timeout_for_rate_limit_error_with_fallbacks():
-    """
-    Test 2. Timeout is 0.0 when RateLimit Error and fallbacks are > 0
-    """
-    healthy_deployments = None
-    fallbacks = ["fallback1", "fallback2"]
-
-    router = litellm.Router(
-        model_list=[
-            {
-                "model_name": "gpt-3.5-turbo",
-                "litellm_params": {
-                    "model": "azure/chatgpt-v-2",
-                    "api_key": os.getenv("AZURE_API_KEY"),
-                    "api_version": os.getenv("AZURE_API_VERSION"),
-                    "api_base": os.getenv("AZURE_API_BASE"),
-                },
-            }
-        ]
-    )
-
-    _timeout = router._time_to_sleep_before_retry(
-        e=rate_limit_error,
-        remaining_retries=4,
-        num_retries=4,
-        healthy_deployments=healthy_deployments,
-        fallbacks=fallbacks,
-    )
-
-    print(
-        "timeout=",
-        _timeout,
-        "error is rate_limit_error and there are fallbacks=",
-        fallbacks,
-    )
-
-    assert _timeout == 0.0
-
-
 def test_timeout_for_rate_limit_error_with_no_healthy_deployments():
     """
-    Test 3. Timeout is > 0.0 when RateLimit Error and healthy deployments == 0 and fallbacks == None
+    Test 2. Timeout is > 0.0 when RateLimit Error and healthy deployments == 0
     """
     healthy_deployments = []
-    fallbacks = None
 
     router = litellm.Router(
         model_list=[
@@ -538,7 +493,6 @@ def test_timeout_for_rate_limit_error_with_no_healthy_deployments():
         remaining_retries=4,
         num_retries=4,
         healthy_deployments=healthy_deployments,
-        fallbacks=fallbacks,
     )
 
     print(