diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index cff4cf48fc4d..b3852fc085ba 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -688,8 +688,7 @@ async def new_team( # noqa: PLR0915 }, ) - - if (data.max_budget is not None and user_api_key_dict.user_id is not None): + if data.max_budget is not None and user_api_key_dict.user_id is not None: # Fetch user object to get max_budget user_obj = await get_user_object( user_id=user_api_key_dict.user_id, @@ -699,7 +698,7 @@ async def new_team( # noqa: PLR0915 ) if ( - user_obj is not None + user_obj is not None and user_obj.max_budget is not None and data.max_budget > user_obj.max_budget ): @@ -1278,12 +1277,12 @@ async def update_team( updated_kv["model_id"] = _model_id updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[LiteLLM_TeamTable] = ( - await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore - ) + team_row: Optional[ + LiteLLM_TeamTable + ] = await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore ) if team_row is None or team_row.team_id is None: @@ -1698,14 +1697,16 @@ async def team_member_add( complete_team_data=complete_team_data, ) - updated_team, updated_users, updated_team_memberships = ( - await _add_team_members_to_team( - data=data, - complete_team_data=complete_team_data, - prisma_client=prisma_client, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ) + ( + updated_team, + updated_users, + updated_team_memberships, + ) = await _add_team_members_to_team( + data=data, + complete_team_data=complete_team_data, + prisma_client=prisma_client, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, ) # Check if updated_team is None @@ -2240,10 +2241,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) + team_row_base: Optional[ + BaseModel + ] = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} ) if team_row_base is None: raise Exception @@ -2420,11 +2421,11 @@ async def team_info( ) try: - team_info: Optional[BaseModel] = ( - await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, - ) + team_info: Optional[ + BaseModel + ] = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, ) if team_info is None: raise Exception @@ -3311,15 +3312,21 @@ async def update_team_member_permissions( raise HTTPException(status_code=500, detail={"error": "No db connected"}) ## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN - existing_team_row = await get_team_object( - team_id=data.team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=None, - proxy_logging_obj=proxy_logging_obj, - check_cache_only=False, - check_db_only=True, - ) + try: + existing_team_row = await get_team_object( + team_id=data.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + check_cache_only=False, + check_db_only=True, + ) + except Exception: + raise HTTPException( + status_code=404, + detail={"error": f"Team not found for team_id={data.team_id}"}, + ) if existing_team_row is None: raise HTTPException( status_code=404, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index c9b7e0579042..7d072b30e540 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -1,15 +1,11 @@ -import asyncio -import json import os import sys -from typing import Optional, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from fastapi.testclient import TestClient -from litellm._uuid import uuid sys.path.insert( 0, os.path.abspath("../../../") @@ -28,9 +24,6 @@ user_api_key_auth, # Assuming this dependency is needed ) from litellm.proxy.management_endpoints.team_endpoints import ( - GetTeamMemberPermissionsResponse, - UpdateTeamMemberPermissionsRequest, - router, team_member_add_duplication_check, validate_team_org_change, ) @@ -329,10 +322,10 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): assert mock_team_create.call_count == 1 created_team_kwargs = mock_team_create.call_args.kwargs team_data = created_team_kwargs["data"] - + # Verify object_permission_id is in the team data assert team_data.get("object_permission_id") == "objperm123" - + # Verify object_permission dict is NOT in the team data assert "object_permission" not in team_data @@ -341,7 +334,7 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_auth): """ Test that /team/new correctly handles mcp_tool_permissions in object_permission. - + This test verifies that: 1. mcp_tool_permissions is accepted in the object_permission field 2. The field is properly stored in the LiteLLM_ObjectPermissionTable @@ -379,9 +372,13 @@ async def mock_obj_perm_create(**kwargs): "object_permission_id": "objperm_team_mcp_456", } mock_db_client.db.litellm_teamtable = MagicMock() - mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable.create = AsyncMock( + return_value=team_create_result + ) mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) - mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable.update = AsyncMock( + return_value=team_create_result + ) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() @@ -414,6 +411,7 @@ async def mock_obj_perm_create(**kwargs): # Verify mcp_tool_permissions was stored import json + assert "mcp_tool_permissions" in created_permission_data # mcp_tool_permissions is stored as a JSON string assert json.loads(created_permission_data["mcp_tool_permissions"]) == { @@ -434,7 +432,6 @@ async def test_team_update_object_permissions_existing_permission(monkeypatch): """ from unittest.mock import AsyncMock, MagicMock - import pytest from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_TeamTable from litellm.proxy.management_endpoints.team_endpoints import ( @@ -506,7 +503,6 @@ async def test_team_update_object_permissions_no_existing_permission(monkeypatch """ from unittest.mock import AsyncMock, MagicMock - import pytest from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_TeamTable from litellm.proxy.management_endpoints.team_endpoints import ( @@ -566,7 +562,6 @@ async def test_team_update_object_permissions_missing_permission_record(monkeypa """ from unittest.mock import AsyncMock, MagicMock - import pytest from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_TeamTable from litellm.proxy.management_endpoints.team_endpoints import ( @@ -1094,7 +1089,6 @@ def test_add_new_models_to_team_with_existing_models(): """ Test add_new_models_to_team function with existing models """ - from litellm.proxy._types import SpecialModelNames from litellm.proxy.management_endpoints.team_endpoints import add_new_models_to_team team_obj = MagicMock(spec=LiteLLM_TeamTable) @@ -1144,7 +1138,6 @@ async def test_update_team_team_member_budget_not_passed_to_db(): ) as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget: - # Setup mock prisma client mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = { @@ -1169,7 +1162,12 @@ async def test_update_team_team_member_budget_not_passed_to_db(): # Mock budget upsert to return updated_kv without team_member_budget def mock_upsert_side_effect( - team_table, user_api_key_dict, updated_kv, team_member_budget=None, team_member_rpm_limit=None, team_member_tpm_limit=None + team_table, + user_api_key_dict, + updated_kv, + team_member_budget=None, + team_member_rpm_limit=None, + team_member_tpm_limit=None, ): # Remove team_member_budget from updated_kv as the real function does result_kv = updated_kv.copy() @@ -1274,7 +1272,6 @@ async def test_bulk_team_member_add_success(): from litellm.proxy._types import ( LiteLLM_TeamMembership, LiteLLM_UserTable, - TeamAddMemberResponse, ) from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add @@ -1345,7 +1342,6 @@ async def test_bulk_team_member_add_success(): new_callable=AsyncMock, return_value=mock_team_response, ) as mock_team_member_add: - mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) result = await bulk_team_member_add( @@ -1434,7 +1430,7 @@ async def test_bulk_team_member_add_all_users_flag(): """ Test bulk_team_member_add with all_users flag set to True """ - from litellm.proxy._types import LiteLLM_UserTable, TeamAddMemberResponse + from litellm.proxy._types import TeamAddMemberResponse from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add bulk_request = BulkTeamMemberAddRequest( @@ -1461,7 +1457,6 @@ async def test_bulk_team_member_add_all_users_flag(): new_callable=AsyncMock, return_value=mock_team_response, ) as mock_team_member_add: - # Mock the database find_many call mock_prisma.db.litellm_usertable.find_many = AsyncMock( return_value=mock_db_users @@ -1509,7 +1504,6 @@ async def test_bulk_team_member_add_failure_scenario(): new_callable=AsyncMock, side_effect=Exception("Database connection failed"), ) as mock_team_member_add: - mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) result = await bulk_team_member_add( @@ -1563,7 +1557,7 @@ async def test_list_team_v2_security_check_non_admin_user(): Test that list_team_v2 properly checks route permissions for non-admin users. Non-admin users should only be able to query their own teams. """ - from unittest.mock import AsyncMock, Mock, patch + from unittest.mock import Mock, patch from fastapi import HTTPException, Request @@ -1603,7 +1597,7 @@ async def test_list_team_v2_security_check_non_admin_user_other_user(): Test that list_team_v2 properly checks route permissions for non-admin users trying to query other users' teams. """ - from unittest.mock import AsyncMock, Mock, patch + from unittest.mock import Mock, patch from fastapi import HTTPException, Request @@ -1661,7 +1655,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db - + # Mock user lookup mock_user_object = Mock() mock_user_object.model_dump.return_value = { @@ -1669,7 +1663,7 @@ async def test_list_team_v2_security_check_non_admin_user_own_teams(): "teams": ["team_1", "team_2"], } mock_db.litellm_usertable.find_unique = AsyncMock(return_value=mock_user_object) - + # Mock team lookup mock_teams = [ Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Team 1"}), @@ -1719,7 +1713,7 @@ async def test_list_team_v2_security_check_admin_user(): # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db - + # Mock team lookup mock_teams = [ Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Team 1"}), @@ -1769,19 +1763,25 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a } # Configure DB mocks used by team_member_delete - mock_db_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_team_row) + mock_db_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=mock_team_row + ) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=mock_team_row) # User row to allow removal from user's teams list mock_user_row = MagicMock() mock_user_row.user_id = test_user_id mock_user_row.teams = [test_team_id] - mock_db_client.db.litellm_usertable.find_many = AsyncMock(return_value=[mock_user_row]) + mock_db_client.db.litellm_usertable.find_many = AsyncMock( + return_value=[mock_user_row] + ) mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) # Membership deletion should be called mock_db_client.db.litellm_teammembership = MagicMock() - mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) + mock_db_client.db.litellm_teammembership.delete_many = AsyncMock( + return_value=MagicMock() + ) # Execute await team_member_delete( @@ -1793,13 +1793,13 @@ async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_a mock_db_client.db.litellm_teammembership.delete_many.assert_awaited_with( where={"team_id": test_team_id, "user_id": test_user_id} ) - + @pytest.mark.asyncio async def test_new_team_max_budget_exceeds_user_max_budget(): """ Test that /team/new raises ProxyException when max_budget exceeds user's end_user_max_budget. - + This validates the budget enforcement logic where non-admin users cannot create teams with budgets higher than their personal maximum budget limit. """ @@ -1836,15 +1836,16 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False mock_prisma.get_data = AsyncMock(return_value=None) - + # Mock user cache to return a user object with max_budget=100.0 from litellm.proxy._types import LiteLLM_UserTable + mock_user_obj = LiteLLM_UserTable( user_id="non-admin-user-123", max_budget=100.0, ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) - + # Should raise ProxyException (HTTPException gets converted by handle_exception_on_proxy) with pytest.raises(ProxyException) as exc_info: await new_team( @@ -1855,9 +1856,11 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): # Verify exception details # ProxyException stores status_code in 'code' attribute - assert exc_info.value.code == '400' + assert exc_info.value.code == "400" assert "max budget higher than user max" in str(exc_info.value.message) - assert "100.0" in str(exc_info.value.message) # User's user_max_budget should be mentioned + assert "100.0" in str( + exc_info.value.message + ) # User's user_max_budget should be mentioned assert LitellmUserRoles.INTERNAL_USER.value in str(exc_info.value.message) @@ -1865,7 +1868,7 @@ async def test_new_team_max_budget_exceeds_user_max_budget(): async def test_new_team_max_budget_within_user_limit(): """ Test that /team/new succeeds when max_budget is within user's user_max_budget. - + This ensures that users can create teams with budgets at or below their personal limit. """ from fastapi import Request @@ -1898,22 +1901,22 @@ async def test_new_team_max_budget_within_user_limit(): ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: - # Setup mocks mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_license.is_team_count_over_limit.return_value = False mock_prisma.jsonify_team_object = lambda db_data: db_data mock_prisma.get_data = AsyncMock(return_value=None) mock_prisma.update_data = AsyncMock() - + # Mock user cache to return a user object with max_budget=100.0 from litellm.proxy._types import LiteLLM_UserTable + mock_user_obj = LiteLLM_UserTable( user_id="non-admin-user-456", max_budget=100.0, ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) - + # Mock team creation mock_created_team = MagicMock() mock_created_team.team_id = "team-within-budget-789" @@ -1927,21 +1930,30 @@ async def test_new_team_max_budget_within_user_limit(): "max_budget": 50.0, "members_with_roles": [], } - mock_prisma.db.litellm_teamtable.create = AsyncMock(return_value=mock_created_team) - mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_created_team) - + mock_prisma.db.litellm_teamtable.create = AsyncMock( + return_value=mock_created_team + ) + mock_prisma.db.litellm_teamtable.update = AsyncMock( + return_value=mock_created_team + ) + # Mock model table mock_prisma.db.litellm_modeltable = MagicMock() - mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) - + mock_prisma.db.litellm_modeltable.create = AsyncMock( + return_value=MagicMock(id="model123") + ) + # Mock user table operations for adding the creator as a member mock_user = MagicMock() mock_user.user_id = "non-admin-user-456" - mock_user.model_dump.return_value = {"user_id": "non-admin-user-456", "teams": ["team-within-budget-789"]} + mock_user.model_dump.return_value = { + "user_id": "non-admin-user-456", + "teams": ["team-within-budget-789"], + } mock_prisma.db.litellm_usertable = MagicMock() mock_prisma.db.litellm_usertable.upsert = AsyncMock(return_value=mock_user) mock_prisma.db.litellm_usertable.update = AsyncMock(return_value=mock_user) - + # Mock team membership table mock_membership = MagicMock() mock_membership.model_dump.return_value = { @@ -1950,7 +1962,9 @@ async def test_new_team_max_budget_within_user_limit(): "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() - mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) + mock_prisma.db.litellm_teammembership.create = AsyncMock( + return_value=mock_membership + ) # Should NOT raise an exception result = await new_team( @@ -1963,3 +1977,37 @@ async def test_new_team_max_budget_within_user_limit(): assert result is not None assert result["team_id"] == "team-within-budget-789" assert result["max_budget"] == 50.0 + + +@pytest.mark.asyncio +async def test_update_team_permissions_fail_team_not_found( + mock_db_client, mock_admin_auth +): + """ + Test update team permissions fails with 404 when team doesn't exist. + """ + test_team_id = "non-existent-team" + update_payload = { + "team_id": test_team_id, + "team_member_permissions": [], + } + + # Mock get_team_object to raise the exception that happens when team is missing + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + side_effect=Exception( + f"Team doesn't exist in db. Team={test_team_id}. Create team via `/team/new` call." + ), + ): + # Override the dependency for this test + app.dependency_overrides[user_api_key_auth] = lambda: mock_admin_auth + + response = client.post("/team/permissions_update", json=update_payload) + + # Assert 404 Not Found instead of 500 + assert response.status_code == 404 + assert "Team not found" in str(response.json()["detail"]) + + # Clean up dependency override + app.dependency_overrides = {}