From 5a2c9d51211746a2d32f8a24c77f5ec9d7362661 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Sat, 24 Aug 2024 10:08:14 -0700
Subject: [PATCH] test(test_router.py): add test to ensure error is correctly
 re-raised

---
 litellm/main.py              |  1 -
 litellm/tests/test_router.py | 53 ++++++++++++++++++++++++++++--------
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/litellm/main.py b/litellm/main.py
index 001cbe24f..649fc12a5 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -503,7 +503,6 @@ def mock_completion(
         ):
             raise litellm.RateLimitError(
                 message="this is a mock rate limit error",
-                status_code=getattr(mock_response, "status_code", 429),  # type: ignore
                 llm_provider=getattr(mock_response, "llm_provider", custom_llm_provider or "openai"),  # type: ignore
                 model=model,
             )
diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py
index 0eb014f85..3c374df87 100644
--- a/litellm/tests/test_router.py
+++ b/litellm/tests/test_router.py
@@ -2110,18 +2110,18 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider):
 
 
 def test_router_cooldown_api_connection_error():
-    # try:
-    #     _ = litellm.completion(
-    #         model="vertex_ai/gemini-1.5-pro",
-    #         messages=[{"role": "admin", "content": "Fail on this!"}],
-    #     )
-    # except litellm.APIConnectionError as e:
-    #     assert (
-    #         Router()._is_cooldown_required(
-    #             exception_status=e.code, exception_str=str(e)
-    #         )
-    #         is False
-    #     )
+    try:
+        _ = litellm.completion(
+            model="vertex_ai/gemini-1.5-pro",
+            messages=[{"role": "admin", "content": "Fail on this!"}],
+        )
+    except litellm.APIConnectionError as e:
+        assert (
+            Router()._is_cooldown_required(
+                exception_status=e.code, exception_str=str(e)
+            )
+            is False
+        )
 
     router = Router(
         model_list=[
@@ -2155,3 +2155,32 @@ def test_router_cooldown_api_connection_error():
         )
     except litellm.APIConnectionError:
         pass
+
+
+def test_router_correctly_reraise_error():
+    """
+    User feedback: There is a problem with my messages array, but the error exception thrown is a Rate Limit error.
+    ```
+    Rate Limit: Error code: 429 - {'error': {'message': 'No deployments available for selected model, Try again in 60 seconds. Passed model=gemini-1.5-flash..
+    ```
+    What they want? Propagation of the real error.
+    """
+    router = Router(
+        model_list=[
+            {
+                "model_name": "gemini-1.5-pro",
+                "litellm_params": {
+                    "model": "vertex_ai/gemini-1.5-pro",
+                    "mock_response": "litellm.RateLimitError",
+                },
+            }
+        ]
+    )
+
+    try:
+        router.completion(
+            model="gemini-1.5-pro",
+            messages=[{"role": "admin", "content": "Fail on this!"}],
+        )
+    except litellm.RateLimitError:
+        pass