From fdb47e54795267227923e1c28f08596a5143a060 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 5 Aug 2024 21:39:45 -0700
Subject: [PATCH] fix: fix test to specify allowed_fails

---
 litellm/router.py                       |  5 ++++-
 litellm/tests/test_acooldowns_router.py | 18 +++++++++++++-----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index e31de5332e..aa9768ba44 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -321,7 +321,10 @@ class Router:
                 []
             )  # initialize an empty list - to allow _add_deployment and delete_deployment to work
 
-        self.allowed_fails = allowed_fails or litellm.allowed_fails
+        if allowed_fails is not None:
+            self.allowed_fails = allowed_fails
+        else:
+            self.allowed_fails = litellm.allowed_fails
         self.cooldown_time = cooldown_time or 60
         self.disable_cooldowns = disable_cooldowns
         self.failed_calls = (
diff --git a/litellm/tests/test_acooldowns_router.py b/litellm/tests/test_acooldowns_router.py
index 7dced31a8f..f07bb2c667 100644
--- a/litellm/tests/test_acooldowns_router.py
+++ b/litellm/tests/test_acooldowns_router.py
@@ -1,17 +1,23 @@
 #### What this tests ####
 #    This tests calling batch_completions by running 100 messages together
 
-import sys, os, time
-import traceback, asyncio
+import asyncio
+import os
+import sys
+import time
+import traceback
+
 import pytest
 
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
+import concurrent
+
+from dotenv import load_dotenv
+
 import litellm
 from litellm import Router
-import concurrent
-from dotenv import load_dotenv
 
 load_dotenv()
 
@@ -45,7 +51,8 @@ kwargs = {
 
 
 def test_multiple_deployments_sync():
-    import concurrent, time
+    import concurrent
+    import time
 
     litellm.set_verbose = False
     results = []
@@ -161,6 +168,7 @@ async def test_cooldown_same_model_name(sync_mode):
             routing_strategy="simple-shuffle",
             set_verbose=True,
             num_retries=3,
+            allowed_fails=0,
         )  # type: ignore
 
         if sync_mode: