diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 721a8c2a1c4..8f6b8a626e4 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -33,7 +33,6 @@ from litellm.proxy._types import ( ProxyException, TeamModelAddRequest, TeamModelDeleteRequest, - UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -42,7 +41,6 @@ from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( team_model_add, team_model_delete, - update_team, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.utils import PrismaClient diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py index 90d98b8ab0a..2694c62827c 100644 --- a/tests/router_unit_tests/test_router_index_management.py +++ b/tests/router_unit_tests/test_router_index_management.py @@ -118,6 +118,28 @@ class TestRouterIndexManagement: assert router.model_id_to_deployment_index_map["id-2"] == 1 assert router.model_id_to_deployment_index_map["id-3"] == 2 + def test_update_team_model_index(self, router): + """Test _update_team_model_index updates team_model_to_deployment_indices.""" + model = { + "model_name": "team-alias", + "model_info": { + "id": "dep-1", + "team_id": "team-abc", + "team_public_model_name": "gpt-4o", + }, + } + router._update_team_model_index(model, 0) + assert router.team_model_to_deployment_indices[("team-abc", "gpt-4o")] == [0] + router._update_team_model_index(model, 2) + assert router.team_model_to_deployment_indices[("team-abc", "gpt-4o")] == [0, 2] + + router._update_team_model_index( + {"model_name": "x", "model_info": {"id": "dep-2"}}, 5 + ) + assert router.team_model_to_deployment_indices == { + ("team-abc", "gpt-4o"): [0, 2], + } + def test_has_model_id(self, router): """Test has_model_id function for O(1) membership check""" # Setup: Add models to router