fix team endpoints test

tests/litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -59,40 +59,47 @@ async def test_get_team_permissions_list_success(mock_db_client, mock_admin_auth
     Test successful retrieval of team member permissions.
     """
     test_team_id = "test-team-123"
+    permissions = ["/key/generate", "/key/update"]
     mock_team_data = {
         "team_id": test_team_id,
         "team_alias": "Test Team",
-        "team_member_permissions": ["/key/generate", "/key/update"],
+        "team_member_permissions": permissions,
         "spend": 0.0,
     }
     mock_team_row = MagicMock()
     mock_team_row.model_dump.return_value = mock_team_data
-    mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(
-        return_value=mock_team_row
-    )
 
-    # Override the dependency for this test
-    app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth
+    # Set attributes directly on the mock object
+    mock_team_row.team_id = test_team_id
+    mock_team_row.team_alias = "Test Team"
+    mock_team_row.team_member_permissions = permissions
+    mock_team_row.spend = 0.0
 
-    response = client.get(f"/team/permissions_list?team_id={test_team_id}")
+    # Mock the get_team_object function used in the endpoint
+    with patch(
+        "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+        new_callable=AsyncMock,
+        return_value=mock_team_row,
+    ):
+        # Override the dependency for this test
+        app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth
 
-    assert response.status_code == 200
-    response_data = response.json()
-    assert response_data["team_id"] == test_team_id
-    assert (
-        response_data["team_member_permissions"]
-        == mock_team_data["team_member_permissions"]
-    )
-    assert (
-        response_data["all_available_permissions"]
-        == TeamMemberPermissionChecks.get_all_available_team_member_permissions()
-    )
-    mock_db_client.db.litellm_teamtable.find_unique.assert_awaited_once_with(
-        where={"team_id": test_team_id}
-    )
+        response = client.get(f"/team/permissions_list?team_id={test_team_id}")
 
-    # Clean up dependency override
-    app.dependency_overrides = {}
+        assert response.status_code == 200
+        response_data = response.json()
+        assert response_data["team_id"] == test_team_id
+        assert (
+            response_data["team_member_permissions"]
+            == mock_team_data["team_member_permissions"]
+        )
+        assert (
+            response_data["all_available_permissions"]
+            == TeamMemberPermissionChecks.get_all_available_team_member_permissions()
+        )
+
+        # Clean up dependency override
+        app.dependency_overrides = {}
 
 
 # Test for /team/permissions_update endpoint (POST)
@@ -102,15 +109,17 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth):
     Test successful update of team member permissions by an admin.
     """
     test_team_id = "test-team-456"
+    update_permissions = ["/key/generate", "/key/update"]
     update_payload = {
         "team_id": test_team_id,
-        "team_member_permissions": ["/key/generate", "/key/update"],
+        "team_member_permissions": update_permissions,
     }
 
+    existing_permissions = ["/key/list"]
     mock_existing_team_data = {
         "team_id": test_team_id,
         "team_alias": "Existing Team",
-        "team_member_permissions": ["/key/list"],
+        "team_member_permissions": existing_permissions,
         "spend": 0.0,
         "models": [],
     }
@@ -121,41 +130,50 @@ async def test_update_team_permissions_success(mock_db_client, mock_admin_auth):
 
     mock_existing_team_row = MagicMock(spec=LiteLLM_TeamTable)
     mock_existing_team_row.model_dump.return_value = mock_existing_team_data
-    # Set attributes directly if model_dump isn't enough for LiteLLM_TeamTable usage
-    for key, value in mock_existing_team_data.items():
-        setattr(mock_existing_team_row, key, value)
+
+    # Set attributes directly on the existing team mock
+    mock_existing_team_row.team_id = test_team_id
+    mock_existing_team_row.team_alias = "Existing Team"
+    mock_existing_team_row.team_member_permissions = existing_permissions
+    mock_existing_team_row.spend = 0.0
+    mock_existing_team_row.models = []
 
     mock_updated_team_row = MagicMock(spec=LiteLLM_TeamTable)
     mock_updated_team_row.model_dump.return_value = mock_updated_team_data
-    # Set attributes directly if model_dump isn't enough for LiteLLM_TeamTable usage
-    for key, value in mock_updated_team_data.items():
-        setattr(mock_updated_team_row, key, value)
 
-    mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(
-        return_value=mock_existing_team_row
-    )
-    mock_db_client.db.litellm_teamtable.update = AsyncMock(
-        return_value=mock_updated_team_row
-    )
+    # Set attributes directly on the updated team mock
+    mock_updated_team_row.team_id = test_team_id
+    mock_updated_team_row.team_alias = "Existing Team"
+    mock_updated_team_row.team_member_permissions = update_permissions
+    mock_updated_team_row.spend = 0.0
+    mock_updated_team_row.models = []
 
-    # Override the dependency for this test
-    app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth
+    # Mock the get_team_object function used in the endpoint
+    with patch(
+        "litellm.proxy.management_endpoints.team_endpoints.get_team_object",
+        new_callable=AsyncMock,
+        return_value=mock_existing_team_row,
+    ):
+        # Mock the database update function
+        mock_db_client.db.litellm_teamtable.update = AsyncMock(
+            return_value=mock_updated_team_row
+        )
 
-    response = client.post("/team/permissions_update", json=update_payload)
+        # Override the dependency for this test
+        app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth
 
-    assert response.status_code == 200
-    response_data = response.json()
+        response = client.post("/team/permissions_update", json=update_payload)
 
-    # Use model_dump for comparison if the endpoint returns the Prisma model directly
-    assert response_data == mock_updated_team_row.model_dump()
+        assert response.status_code == 200
+        response_data = response.json()
 
-    mock_db_client.db.litellm_teamtable.find_unique.assert_awaited_once_with(
-        where={"team_id": test_team_id}
-    )
-    mock_db_client.db.litellm_teamtable.update.assert_awaited_once_with(
-        where={"team_id": test_team_id},
-        data={"team_member_permissions": update_payload["team_member_permissions"]},
-    )
+        # Use model_dump for comparison if the endpoint returns the Prisma model directly
+        assert response_data == mock_updated_team_row.model_dump()
 
-    # Clean up dependency override
-    app.dependency_overrides = {}
+        mock_db_client.db.litellm_teamtable.update.assert_awaited_once_with(
+            where={"team_id": test_team_id},
+            data={"team_member_permissions": update_payload["team_member_permissions"]},
+        )
+
+        # Clean up dependency override
+        app.dependency_overrides = {}