From 8687ae800d9c2133b9ef4cbda8a447bfadd4fa7d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 14 Sep 2024 11:49:18 -0700 Subject: [PATCH] fix test_single_deployment_no_cooldowns_test_prod_mock_completion_calls --- litellm/tests/test_router_cooldowns.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_router_cooldowns.py b/litellm/tests/test_router_cooldowns.py index 4287659a35f..b7993d8a539 100644 --- a/litellm/tests/test_router_cooldowns.py +++ b/litellm/tests/test_router_cooldowns.py @@ -21,6 +21,7 @@ import openai import litellm from litellm import Router from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict @@ -239,7 +240,9 @@ async def test_single_deployment_no_cooldowns_test_prod_mock_completion_calls(): except litellm.RateLimitError: pass - cooldown_list = await router._async_get_cooldown_deployments() + cooldown_list = await _async_get_cooldown_deployments( + litellm_router_instance=router + ) assert len(cooldown_list) == 0 healthy_deployments, _ = await router._async_get_healthy_deployments(