From e41b3bd13fa6415de7b4076dda82f673fac8b957 Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 13 Sep 2026 09:50:12 +0000 Subject: [PATCH] test(router): annotate return types of team cooldown test helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/router_utils/test_cooldown_handlers.py | 6 ++---- tests/test_litellm/test_router.py | 4 +++- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/router_utils/test_cooldown_handlers.py b/tests/test_litellm/router_utils/test_cooldown_handlers.py index fdbfd618dab..6fed4be5909 100644 --- a/tests/test_litellm/router_utils/test_cooldown_handlers.py +++ b/tests/test_litellm/router_utils/test_cooldown_handlers.py @@ -440,10 +440,8 @@ class TestRoutingGroupCooldownAlternatives: class TestTeamModelCooldownAlternatives: - def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()): - from litellm import Router - - return Router( + def _router(self, team_deployments: int, blocked_ids: frozenset[str] = frozenset()) -> litellm.Router: + return litellm.Router( model_list=[ { "model_name": f"model_name_team-1_{i}", diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 7786bfb4788..c6b3b7fb8d6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -897,7 +897,9 @@ def test_arouter_test_team_model(): def test_team_model_has_alternatives(): - def team_deployment(deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False): + def team_deployment( + deployment_id: str, team_id: str, public_model_name: str, blocked: bool = False + ) -> DeploymentTypedDict: return { "model_name": f"model_name_{team_id}_{deployment_id}", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},