Fix code qa issues

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

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