From beac60ed12b81d4525000e6f69918caed50a4531 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 11 May 2024 19:58:17 -0700
Subject: [PATCH] test - router retry policy

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

diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py
index 3a89b644b5..7273fd6e96 100644
--- a/litellm/tests/test_router_retries.py
+++ b/litellm/tests/test_router_retries.py
@@ -192,8 +192,8 @@ async def test_dynamic_router_retry_policy(model_group):
     from litellm.router import RetryPolicy
 
     model_group_retry_policy = {
-        "gpt-3.5-turbo": RetryPolicy(ContentPolicyViolationErrorRetries=0),
-        "bad-model": RetryPolicy(AuthenticationErrorRetries=4),
+        "gpt-3.5-turbo": RetryPolicy(ContentPolicyViolationErrorRetries=2),
+        "bad-model": RetryPolicy(AuthenticationErrorRetries=0),
     }
 
     router = litellm.Router(
@@ -206,6 +206,33 @@ async def test_dynamic_router_retry_policy(model_group):
                     "api_version": os.getenv("AZURE_API_VERSION"),
                     "api_base": os.getenv("AZURE_API_BASE"),
                 },
+                "model_info": {
+                    "id": "model-0",
+                },
+            },
+            {
+                "model_name": "gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "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"),
+                },
+                "model_info": {
+                    "id": "model-1",
+                },
+            },
+            {
+                "model_name": "gpt-3.5-turbo",  # openai model name
+                "litellm_params": {  # params for litellm completion/embedding call
+                    "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"),
+                },
+                "model_info": {
+                    "id": "model-2",
+                },
             },
             {
                 "model_name": "bad-model",  # openai model name
@@ -241,9 +268,9 @@ async def test_dynamic_router_retry_policy(model_group):
     print("customHandler.previous_models: ", customHandler.previous_models)
 
     if model_group == "bad-model":
-        assert customHandler.previous_models == 4
-    elif model_group == "gpt-3.5-turbo":
         assert customHandler.previous_models == 0
+    elif model_group == "gpt-3.5-turbo":
+        assert customHandler.previous_models == 2
 
 
 """