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("../../../") ) # Adds the parent directory to the system path from litellm.proxy._types import UserAPIKeyAuth # Import UserAPIKeyAuth from litellm.proxy._types import ( LiteLLM_OrganizationMembershipTable, LiteLLM_OrganizationTable, LiteLLM_OrganizationTableWithMembers, LiteLLM_TeamTable, LiteLLM_UserTable, LitellmUserRoles, Member, ProxyErrorTypes, ProxyException, TeamMemberAddRequest, ) from litellm.proxy.management_endpoints.team_endpoints import ( user_api_key_auth, # Assuming this dependency is needed ) from litellm.proxy.management_endpoints.team_endpoints import ( GetTeamMemberPermissionsResponse, UpdateTeamMemberPermissionsRequest, _persist_deleted_team_records, _save_deleted_team_records, _transform_teams_to_deleted_records, _validate_and_populate_member_user_info, delete_team, list_available_teams, router, team_member_add_duplication_check, team_member_delete, validate_team_org_change, ) from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) from litellm.proxy.proxy_server import app from litellm.router import Router from litellm.types.proxy.management_endpoints.team_endpoints import ( BulkTeamMemberAddRequest, BulkTeamMemberAddResponse, TeamMemberAddResult, ) # Setup TestClient client = TestClient(app) # Mock prisma_client mock_prisma_client = MagicMock() # Set up async mock for db operations mock_prisma_client.db = MagicMock() mock_prisma_client.db.litellm_teamtable = MagicMock() mock_prisma_client.db.litellm_teamtable.update = AsyncMock() # Fixture to provide the mock prisma client @pytest.fixture(autouse=True) def mock_db_client(): with patch( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): # Mock in both places if necessary yield mock_prisma_client mock_prisma_client.reset_mock() # Fixture to provide a mock admin user auth object @pytest.fixture def mock_admin_auth(): mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) return mock_auth # Test for validate_team_org_change when organization IDs match @pytest.mark.asyncio async def test_validate_team_org_change_same_org_id(): """ Test that validate_team_org_change returns True without performing any checks when the team and organization have the same organization_id. This is a user issue, a user was editing their team and this function raised an exception even when they were not changing the organization. """ # Create mock team and organization with same org ID org_id = "test-org-123" # Mock team team = MagicMock(spec=LiteLLM_TeamTable) team.organization_id = org_id team.models = ["gpt-4", "claude-2"] team.max_budget = 100.0 team.tpm_limit = 1000 team.rpm_limit = 100 team.members_with_roles = [] # Mock organization organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) organization.organization_id = org_id organization.models = [] organization.litellm_budget_table = MagicMock() organization.litellm_budget_table.max_budget = ( 50.0 # This would normally fail validation ) organization.litellm_budget_table.tpm_limit = ( 500 # This would normally fail validation ) organization.litellm_budget_table.rpm_limit = ( 50 # This would normally fail validation ) organization.members = [] # Mock Router mock_router = MagicMock(spec=Router) # Use patch to ensure the model access check is never called with patch( "litellm.proxy.management_endpoints.team_endpoints.can_org_access_model" ) as mock_access_check: result = validate_team_org_change( team=team, organization=organization, llm_router=mock_router ) # Assert the function returns True without checking anything assert result is True mock_access_check.assert_not_called() # Ensure access check wasn't called @pytest.mark.asyncio async def test_validate_team_org_change_members_in_org(): """ Test that validate_team_org_change passes when team members are in organization.members. This tests the fix for issue #17552 where membership was incorrectly checked against organization.users (deprecated) instead of organization.members (correct). """ team_org_id = "team-org-123" new_org_id = "new-org-456" user_id_1 = "user-123" user_id_2 = "user-456" # Mock team with members team = MagicMock(spec=LiteLLM_TeamTable) team.organization_id = team_org_id team.models = [] team.max_budget = None team.tpm_limit = None team.rpm_limit = None # Create mock team members team_member_1 = MagicMock() team_member_1.user_id = user_id_1 team_member_2 = MagicMock() team_member_2.user_id = user_id_2 team.members_with_roles = [team_member_1, team_member_2] # Mock organization with members (using LiteLLM_OrganizationMembershipTable structure) organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) organization.organization_id = new_org_id organization.models = [] organization.litellm_budget_table = None # Create mock organization members - these should match team members org_member_1 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) org_member_1.user_id = user_id_1 org_member_2 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) org_member_2.user_id = user_id_2 organization.members = [org_member_1, org_member_2] # Mock Router mock_router = MagicMock(spec=Router) # Test should pass - all team members are in org members result = validate_team_org_change( team=team, organization=organization, llm_router=mock_router ) assert result is True @pytest.mark.asyncio async def test_validate_team_org_change_member_not_in_org(): """ Test that validate_team_org_change raises HTTPException when team members are NOT in organization.members. This tests the fix for issue #17552 where membership was incorrectly checked against organization.users (deprecated) instead of organization.members (correct). """ team_org_id = "team-org-123" new_org_id = "new-org-456" user_id_1 = "user-123" user_id_2 = "user-456" user_id_not_in_org = "user-not-in-org-789" # Mock team with members (including one not in org) team = MagicMock(spec=LiteLLM_TeamTable) team.organization_id = team_org_id team.models = [] team.max_budget = None team.tpm_limit = None team.rpm_limit = None # Create mock team members - user_id_not_in_org is not in the org team_member_1 = MagicMock() team_member_1.user_id = user_id_1 team_member_2 = MagicMock() team_member_2.user_id = user_id_not_in_org team.members_with_roles = [team_member_1, team_member_2] # Mock organization with members (missing user_id_not_in_org) organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers) organization.organization_id = new_org_id organization.models = [] organization.litellm_budget_table = None # Create mock organization members - only user_id_1 and user_id_2 are members org_member_1 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) org_member_1.user_id = user_id_1 org_member_2 = MagicMock(spec=LiteLLM_OrganizationMembershipTable) org_member_2.user_id = user_id_2 organization.members = [org_member_1, org_member_2] # Mock Router mock_router = MagicMock(spec=Router) # Test should fail - user_id_not_in_org is not in org members with pytest.raises(HTTPException) as exc_info: validate_team_org_change( team=team, organization=organization, llm_router=mock_router ) assert exc_info.value.status_code == 403 assert "not a member of the organization" in str(exc_info.value.detail) assert user_id_not_in_org in str(exc_info.value.detail) # Test for /team/permissions_list endpoint (GET) @pytest.mark.asyncio 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": permissions, "spend": 0.0, } mock_team_row = MagicMock() mock_team_row.model_dump.return_value = mock_team_data # 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 # 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 response = client.get(f"/team/permissions_list?team_id={test_team_id}") 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) @pytest.mark.asyncio 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": update_permissions, } existing_permissions = ["/key/list"] mock_existing_team_data = { "team_id": test_team_id, "team_alias": "Existing Team", "team_member_permissions": existing_permissions, "spend": 0.0, "models": [], } mock_updated_team_data = { **mock_existing_team_data, "team_member_permissions": update_payload["team_member_permissions"], } mock_existing_team_row = MagicMock(spec=LiteLLM_TeamTable) mock_existing_team_row.model_dump.return_value = mock_existing_team_data # 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 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 = [] # 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 ) # 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 response.status_code == 200 response_data = response.json() # Use model_dump for comparison if the endpoint returns the Prisma model directly assert response_data == mock_updated_team_row.model_dump() 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 = {} @pytest.mark.asyncio async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): """ Test that /team/new correctly handles object_permission by: 1. Creating a record in litellm_objectpermissiontable 2. Passing the returned object_permission_id into the team insert payload 3. NOT passing the object_permission dict to the team table """ # Configure mocked prisma client mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.get_data = AsyncMock(return_value=None) mock_db_client.update_data = AsyncMock(return_value=MagicMock()) mock_db_client.db = MagicMock() # Mock object permission table creation mock_object_perm_create = AsyncMock( return_value=MagicMock(object_permission_id="objperm123") ) mock_db_client.db.litellm_objectpermissiontable = MagicMock() mock_db_client.db.litellm_objectpermissiontable.create = mock_object_perm_create # Mock model table creation mock_db_client.db.litellm_modeltable = MagicMock() mock_db_client.db.litellm_modeltable.create = AsyncMock( return_value=MagicMock(id="model123") ) # Capture team table creation team_create_result = MagicMock( team_id="team-456", object_permission_id="objperm123", ) team_create_result.model_dump.return_value = { "team_id": "team-456", "object_permission_id": "objperm123", } mock_team_create = AsyncMock(return_value=team_create_result) mock_team_count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create mock_db_client.db.litellm_teamtable.count = mock_team_count mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result ) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) from fastapi import Request from litellm.proxy._types import LiteLLM_ObjectPermissionBase, NewTeamRequest from litellm.proxy.management_endpoints.team_endpoints import new_team # Build request with object_permission team_request = NewTeamRequest( team_alias="my-team", object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["my-vector"]), ) dummy_request = MagicMock(spec=Request) # Execute the endpoint function await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=mock_admin_auth, ) # Verify object permission creation was called mock_object_perm_create.assert_awaited_once() # Verify team creation was called 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 @pytest.mark.asyncio 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 3. The team is correctly linked to the object_permission record """ # Configure mocked prisma client mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.get_data = AsyncMock(return_value=None) mock_db_client.update_data = AsyncMock(return_value=MagicMock()) mock_db_client.db = MagicMock() # Track what data is passed to object permission create created_permission_data = {} async def mock_obj_perm_create(**kwargs): created_permission_data.update(kwargs.get("data", {})) return MagicMock(object_permission_id="objperm_team_mcp_456") mock_db_client.db.litellm_objectpermissiontable = MagicMock() mock_db_client.db.litellm_objectpermissiontable.create = mock_obj_perm_create # Mock model table mock_db_client.db.litellm_modeltable = MagicMock() mock_db_client.db.litellm_modeltable.create = AsyncMock( return_value=MagicMock(id="model456") ) # Mock team table team_create_result = MagicMock( team_id="team-mcp-789", object_permission_id="objperm_team_mcp_456", ) team_create_result.model_dump.return_value = { "team_id": "team-mcp-789", "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.count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) from fastapi import Request from litellm.proxy._types import LiteLLM_ObjectPermissionBase, NewTeamRequest from litellm.proxy.management_endpoints.team_endpoints import new_team # Create team with mcp_tool_permissions team_request = NewTeamRequest( team_alias="mcp-team", object_permission=LiteLLM_ObjectPermissionBase( mcp_servers=["server_a", "server_b"], mcp_tool_permissions={ "server_a": ["read_wiki_structure", "read_wiki_contents"], "server_b": ["ask_question"], }, ), ) dummy_request = MagicMock(spec=Request) await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=mock_admin_auth, ) # 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"]) == { "server_a": ["read_wiki_structure", "read_wiki_contents"], "server_b": ["ask_question"], } assert created_permission_data["mcp_servers"] == ["server_a", "server_b"] @pytest.mark.asyncio async def test_team_update_object_permissions_existing_permission(monkeypatch): """ Test updating object permissions when a team already has an existing object_permission_id. This test verifies that when updating vector stores for a team that already has an object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated with the new permissions and the object_permission_id remains the same. """ 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 ( handle_update_object_permission, ) # Mock prisma client mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Mock existing team with object_permission_id existing_team_row = LiteLLM_TeamTable( team_id="test_team_id", object_permission_id="existing_perm_id_123", team_alias="test_team", ) # Mock existing object permission record existing_object_permission = MagicMock() existing_object_permission.model_dump.return_value = { "object_permission_id": "existing_perm_id_123", "vector_stores": ["old_store_1", "old_store_2"], } mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( return_value=existing_object_permission ) # Mock upsert operation updated_permission = MagicMock() updated_permission.object_permission_id = "existing_perm_id_123" mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( return_value=updated_permission ) # Test data with new object permission data_json = { "object_permission": LiteLLM_ObjectPermissionBase( vector_stores=["new_store_1", "new_store_2", "new_store_3"] ).model_dump(exclude_unset=True, exclude_none=True), "team_alias": "updated_team", } # Call the function result = await handle_update_object_permission( data_json=data_json, existing_team_row=existing_team_row, ) # Verify the object_permission was removed from data_json and object_permission_id was set assert "object_permission" not in result assert result["object_permission_id"] == "existing_perm_id_123" # Verify database operations were called correctly mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with( where={"object_permission_id": "existing_perm_id_123"} ) mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() @pytest.mark.asyncio async def test_team_update_object_permissions_no_existing_permission(monkeypatch): """ Test creating object permissions when a team has no existing object_permission_id. This test verifies that when updating object permissions for a team that has object_permission_id set to None, a new entry is created in the LiteLLM_ObjectPermissionTable and the team is updated with the new object_permission_id. """ 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 ( handle_update_object_permission, ) # Mock prisma client mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_team_row_no_perm = LiteLLM_TeamTable( team_id="test_team_id_2", object_permission_id=None, team_alias="test_team_2", ) # Mock find_unique to return None (no existing permission) mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( return_value=None ) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "new_perm_id_456" mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( return_value=new_permission ) data_json = { "object_permission": LiteLLM_ObjectPermissionBase( vector_stores=["brand_new_store"] ).model_dump(exclude_unset=True, exclude_none=True), "team_alias": "updated_team_2", } result = await handle_update_object_permission( data_json=data_json, existing_team_row=existing_team_row_no_perm, ) # Verify new object_permission_id was set assert "object_permission" not in result assert result["object_permission_id"] == "new_perm_id_456" # Verify upsert was called to create new record mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() @pytest.mark.asyncio async def test_team_update_object_permissions_missing_permission_record(monkeypatch): """ Test creating object permissions when existing object_permission_id record is not found. This test verifies that when updating object permissions for a team that has an object_permission_id but the corresponding record cannot be found in the database, a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions. """ 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 ( handle_update_object_permission, ) # Mock prisma client mock_prisma_client = AsyncMock() monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) existing_team_row_missing_perm = LiteLLM_TeamTable( team_id="test_team_id_3", object_permission_id="missing_perm_id_789", team_alias="test_team_3", ) # Mock find_unique to return None (permission record not found) mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock( return_value=None ) # Mock upsert to create new record new_permission = MagicMock() new_permission.object_permission_id = "recreated_perm_id_789" mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock( return_value=new_permission ) data_json = { "object_permission": LiteLLM_ObjectPermissionBase( vector_stores=["recreated_store"] ).model_dump(exclude_unset=True, exclude_none=True), "team_alias": "updated_team_3", } result = await handle_update_object_permission( data_json=data_json, existing_team_row=existing_team_row_missing_perm, ) # Verify new object_permission_id was set assert "object_permission" not in result assert result["object_permission_id"] == "recreated_perm_id_789" # Verify find_unique was called with the missing permission ID mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with( where={"object_permission_id": "missing_perm_id_789"} ) # Verify upsert was called to create new record mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() def test_team_member_add_duplication_check_raises_proxy_exception(): """ Test that team_member_add_duplication_check raises ProxyException when a user is already in the team """ # Create a mock team with existing members existing_team_row = MagicMock(spec=LiteLLM_TeamTable) existing_team_row.team_id = "test-team-123" existing_team_row.members_with_roles = [ Member(user_id="existing-user-id", role="user"), Member(user_id="another-user-id", role="admin"), ] # Create a request to add a member who is already in the team duplicate_member = Member(user_id="existing-user-id", role="user") data = TeamMemberAddRequest( team_id="test-team-123", member=duplicate_member, ) # Test that ProxyException is raised with the correct error type with pytest.raises(ProxyException) as exc_info: team_member_add_duplication_check( data=data, existing_team_row=existing_team_row, ) # Verify the exception details assert exc_info.value.type == ProxyErrorTypes.team_member_already_in_team assert exc_info.value.param == "member" assert exc_info.value.code == "400" assert "existing-user-id" in str(exc_info.value.message) assert "already in team" in str(exc_info.value.message) def test_team_member_add_duplication_check_allows_new_member(): """ Test that team_member_add_duplication_check allows adding a new member who is not already in the team """ # Create a mock team with existing members existing_team_row = MagicMock(spec=LiteLLM_TeamTable) existing_team_row.team_id = "test-team-123" existing_team_row.members_with_roles = [ Member(user_id="existing-user-id", role="user"), Member(user_id="another-user-id", role="admin"), ] # Create a request to add a member who is NOT already in the team new_member = Member(user_id="new-user-id", role="user") data = TeamMemberAddRequest( team_id="test-team-123", member=new_member, ) # Test that no exception is raised for a new member try: team_member_add_duplication_check( data=data, existing_team_row=existing_team_row, ) # If we reach here, no exception was raised, which is expected assert True except ProxyException: # If a ProxyException is raised, the test should fail pytest.fail("ProxyException should not be raised for a new member") @pytest.mark.asyncio async def test_add_team_member_budget_table_success(): """ Test _add_team_member_budget_table when budget is found successfully """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable from litellm.proxy.management_endpoints.team_endpoints import ( _add_team_member_budget_table, ) # Mock prisma client mock_prisma_client = MagicMock() # Mock budget record mock_budget_record = MagicMock() mock_budget_record.budget_id = "budget-123" mock_budget_record.max_budget = 1000.0 mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( return_value=mock_budget_record ) # Create team info response object team_info_response = TeamInfoResponseObjectTeamTable( team_id="test-team-123", team_alias="Test Team" ) # Call the function result = await _add_team_member_budget_table( team_member_budget_id="budget-123", prisma_client=mock_prisma_client, team_info_response_object=team_info_response, ) # Verify the result assert result == team_info_response assert result.team_member_budget_table == mock_budget_record # Verify database call was made correctly mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": "budget-123"} ) @pytest.mark.asyncio async def test_add_team_member_budget_table_exception_handling(): """ Test _add_team_member_budget_table when an exception occurs during budget lookup """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable from litellm.proxy.management_endpoints.team_endpoints import ( _add_team_member_budget_table, ) # Mock prisma client to raise an exception mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock( side_effect=Exception("Database connection failed") ) # Create team info response object team_info_response = TeamInfoResponseObjectTeamTable( team_id="test-team-456", team_alias="Test Team 2" ) # Mock the verbose_proxy_logger to capture log calls with patch( "litellm.proxy.management_endpoints.team_endpoints.verbose_proxy_logger" ) as mock_logger: # Call the function result = await _add_team_member_budget_table( team_member_budget_id="nonexistent-budget-456", prisma_client=mock_prisma_client, team_info_response_object=team_info_response, ) # Verify the result is returned even when exception occurs assert result == team_info_response # Verify team_member_budget_table is not set when exception occurs assert ( not hasattr(result, "team_member_budget_table") or result.team_member_budget_table is None ) # Verify the error was logged mock_logger.info.assert_called_once_with( "Team member budget table not found, passed team_member_budget_id=nonexistent-budget-456" ) # Verify database call was attempted mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": "nonexistent-budget-456"} ) @pytest.mark.asyncio async def test_add_team_member_budget_table_budget_not_found(): """ Test _add_team_member_budget_table when budget record is not found (returns None) """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable from litellm.proxy.management_endpoints.team_endpoints import ( _add_team_member_budget_table, ) # Mock prisma client to return None (budget not found) mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) # Create team info response object team_info_response = TeamInfoResponseObjectTeamTable( team_id="test-team-789", team_alias="Test Team 3" ) # Call the function result = await _add_team_member_budget_table( team_member_budget_id="nonexistent-budget-789", prisma_client=mock_prisma_client, team_info_response_object=team_info_response, ) # Verify the result assert result == team_info_response assert result.team_member_budget_table is None # Verify database call was made correctly mock_prisma_client.db.litellm_budgettable.find_unique.assert_called_once_with( where={"budget_id": "nonexistent-budget-789"} ) def test_add_new_models_to_team(): """ Test add_new_models_to_team function """ 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) team_obj.models = [] new_models = ["model4", "model5"] updated_models = add_new_models_to_team(team_obj=team_obj, new_models=new_models) assert ( updated_models.sort() == [ SpecialModelNames.all_proxy_models.value, "model4", "model5", ].sort() ) @pytest.mark.asyncio async def test_validate_team_member_add_permissions_admin(): """ Test _validate_team_member_add_permissions allows proxy admin """ from litellm.proxy.management_endpoints.team_endpoints import ( _validate_team_member_add_permissions, ) # Create admin user admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) # Create mock team team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = "test-team-123" # Should not raise any exception for admin await _validate_team_member_add_permissions( user_api_key_dict=admin_user, complete_team_data=team, ) @pytest.mark.asyncio async def test_validate_team_member_add_permissions_non_admin(): """ Test _validate_team_member_add_permissions raises exception for non-admin non-team-admin """ from litellm.proxy.management_endpoints.team_endpoints import ( _validate_team_member_add_permissions, ) # Create non-admin user regular_user = UserAPIKeyAuth( user_id="regular-user", user_role=LitellmUserRoles.INTERNAL_USER, team_id="different-team", ) # Create mock team team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = "test-team-123" team.members_with_roles = [] team.organization_id = None # Mock the helper functions to return False with patch( "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", return_value=False, ), patch( "litellm.proxy.management_endpoints.team_endpoints._is_available_team", return_value=False, ): # Should raise HTTPException for non-admin with pytest.raises(HTTPException) as exc_info: await _validate_team_member_add_permissions( user_api_key_dict=regular_user, complete_team_data=team, ) assert exc_info.value.status_code == 403 assert "not proxy admin OR team admin" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_process_team_members_single_member(): """ Test _process_team_members with a single member """ from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_UserTable from litellm.proxy.management_endpoints.team_endpoints import _process_team_members # Mock dependencies mock_prisma_client = MagicMock() mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.metadata = {"team_member_budget_id": "budget-123"} # Mock user and membership objects mock_user = MagicMock(spec=LiteLLM_UserTable) mock_user.user_id = "new-user-123" mock_membership = MagicMock(spec=LiteLLM_TeamMembership) # Create request with single member single_member = Member(user_email="new@example.com", role="user") request_data = TeamMemberAddRequest( team_id="test-team-123", member=single_member, ) with patch( "litellm.proxy.management_endpoints.team_endpoints.add_new_member", new_callable=AsyncMock, return_value=(mock_user, mock_membership), ) as mock_add_member: users, memberships = await _process_team_members( data=request_data, complete_team_data=mock_team, prisma_client=mock_prisma_client, user_api_key_dict=UserAPIKeyAuth(), litellm_proxy_admin_name="admin", ) # Verify results assert len(users) == 1 assert len(memberships) == 1 assert users[0] == mock_user assert memberships[0] == mock_membership # Verify add_new_member was called correctly mock_add_member.assert_called_once_with( new_member=single_member, max_budget_in_team=None, prisma_client=mock_prisma_client, user_api_key_dict=UserAPIKeyAuth(), litellm_proxy_admin_name="admin", team_id="test-team-123", default_team_budget_id="budget-123", ) @pytest.mark.asyncio async def test_process_team_members_multiple_members(): """ Test _process_team_members with multiple members """ from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_UserTable from litellm.proxy.management_endpoints.team_endpoints import _process_team_members # Mock dependencies mock_prisma_client = MagicMock() mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.metadata = None # Create multiple members as dictionaries (they will be converted to Member objects) members = [ Member(user_email="user1@example.com", role="user"), Member(user_email="user2@example.com", role="admin"), ] request_data = TeamMemberAddRequest( team_id="test-team-123", member=members, max_budget_in_team=100.0, ) # Mock different users and memberships for each call mock_users = [MagicMock(spec=LiteLLM_UserTable) for _ in range(2)] mock_memberships = [MagicMock(spec=LiteLLM_TeamMembership) for _ in range(2)] with patch( "litellm.proxy.management_endpoints.team_endpoints.add_new_member", new_callable=AsyncMock, side_effect=[ (mock_users[0], mock_memberships[0]), (mock_users[1], mock_memberships[1]), ], ) as mock_add_member: users, memberships = await _process_team_members( data=request_data, complete_team_data=mock_team, prisma_client=mock_prisma_client, user_api_key_dict=UserAPIKeyAuth(), litellm_proxy_admin_name="admin", ) # Verify results assert len(users) == 2 assert len(memberships) == 2 assert users == mock_users assert memberships == mock_memberships # Verify add_new_member was called for each member assert mock_add_member.call_count == 2 @pytest.mark.asyncio async def test_update_team_members_list_single_member(): """ Test _update_team_members_list with a single member """ from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.management_endpoints.team_endpoints import ( _update_team_members_list, ) # Create mock team with existing members mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.members_with_roles = [Member(user_id="existing-user", role="admin")] # Create new member without user_id new_member = Member(user_email="new@example.com", role="user") request_data = TeamMemberAddRequest( team_id="test-team-123", member=new_member, ) # Create mock user with matching email mock_user = MagicMock(spec=LiteLLM_UserTable) mock_user.user_id = "new-user-123" mock_user.user_email = "new@example.com" await _update_team_members_list( data=request_data, complete_team_data=mock_team, updated_users=[mock_user], ) # Verify member was added assert len(mock_team.members_with_roles) == 2 added_member = mock_team.members_with_roles[1] assert added_member.user_id == "new-user-123" assert added_member.user_email == "new@example.com" assert added_member.role == "user" @pytest.mark.asyncio async def test_update_team_members_list_duplicate_prevention(): """ Test _update_team_members_list prevents duplicate members """ from litellm.proxy._types import LiteLLM_UserTable from litellm.proxy.management_endpoints.team_endpoints import ( _update_team_members_list, ) # Create mock team with existing members mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.members_with_roles = [ Member(user_id="existing-user", user_email="existing@example.com", role="admin") ] # Try to add the same member again duplicate_member = Member(user_id="existing-user", role="user") request_data = TeamMemberAddRequest( team_id="test-team-123", member=duplicate_member, ) # Create mock user mock_user = MagicMock(spec=LiteLLM_UserTable) mock_user.user_id = "existing-user" mock_user.user_email = "existing@example.com" await _update_team_members_list( data=request_data, complete_team_data=mock_team, updated_users=[mock_user], ) # Verify member was NOT added (still only 1 member) assert len(mock_team.members_with_roles) == 1 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) team_obj.models = ["model1", "model2"] new_models = ["model3", "model4"] updated_models = add_new_models_to_team( team_obj=team_obj, new_models=new_models, ) assert updated_models.sort() == ["model1", "model2", "model3", "model4"].sort() @pytest.mark.asyncio async def test_update_team_team_member_budget_not_passed_to_db(): """ Test that 'team_member_budget' is never passed to prisma_client.db.litellm_teamtable.update regardless of whether the value is set or None. This ensures that team_member_budget is properly handled via the separate budget table and not accidentally passed to the team table update operation. """ from unittest.mock import AsyncMock, MagicMock, Mock, patch from fastapi import Request from litellm.proxy._types import LitellmUserRoles, UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team # Mock dependencies mock_request = Mock(spec=Request) mock_user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( "litellm.proxy.proxy_server.llm_router" ) as mock_llm_router, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.proxy_logging_obj" ) as mock_logging, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.auth.auth_checks._cache_team_object" ) 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 = { "team_id": "test_team_id", "team_alias": "test_team", "metadata": {"team_member_budget_id": "budget_123"}, } mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) # Mock the update return value mock_updated_team = MagicMock() mock_updated_team.team_id = "test_team_id" mock_updated_team.model_dump.return_value = {"team_id": "test_team_id"} mock_prisma_client.db.litellm_teamtable.update = AsyncMock( return_value=mock_updated_team ) mock_prisma_client.jsonify_team_object = MagicMock( side_effect=lambda db_data: db_data ) # 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_member_budget_duration=None ): # Remove team_member_budget from updated_kv as the real function does result_kv = updated_kv.copy() result_kv.pop("team_member_budget", None) return result_kv mock_upsert_budget.side_effect = mock_upsert_side_effect # Test Case 1: team_member_budget is set (not None) update_request_with_budget = UpdateTeamRequest( team_id="test_team_id", team_member_budget=100.0, team_alias="updated_alias" ) result = await update_team( data=update_request_with_budget, http_request=mock_request, user_api_key_dict=mock_user_api_key_dict, ) # Verify update was called assert mock_prisma_client.db.litellm_teamtable.update.called # Get the call arguments call_args = mock_prisma_client.db.litellm_teamtable.update.call_args update_data = call_args[1]["data"] # data parameter from the update call # Verify team_member_budget is NOT in the update data assert ( "team_member_budget" not in update_data ), f"team_member_budget should not be in update data, but found: {update_data}" # Verify other fields are present (team_alias should be there) assert "team_alias" in update_data or "team_id" in str( call_args ), "Expected team update fields should be present" # Reset mock for second test mock_prisma_client.db.litellm_teamtable.update.reset_mock() # Test Case 2: team_member_budget is None update_request_without_budget = UpdateTeamRequest( team_id="test_team_id", team_member_budget=None, team_alias="updated_alias_2", ) result = await update_team( data=update_request_without_budget, http_request=mock_request, user_api_key_dict=mock_user_api_key_dict, ) # Verify update was called again assert mock_prisma_client.db.litellm_teamtable.update.called # Get the call arguments for second call call_args = mock_prisma_client.db.litellm_teamtable.update.call_args update_data = call_args[1]["data"] # data parameter from the update call # Verify team_member_budget is NOT in the update data assert ( "team_member_budget" not in update_data ), f"team_member_budget should not be in update data, but found: {update_data}" # Test Case 3: No team_member_budget field at all (excluded from request) mock_prisma_client.db.litellm_teamtable.update.reset_mock() update_request_no_budget_field = UpdateTeamRequest( team_id="test_team_id", team_alias="updated_alias_3", # team_member_budget not specified at all ) result = await update_team( data=update_request_no_budget_field, http_request=mock_request, user_api_key_dict=mock_user_api_key_dict, ) # Verify update was called again assert mock_prisma_client.db.litellm_teamtable.update.called # Get the call arguments for third call call_args = mock_prisma_client.db.litellm_teamtable.update.call_args update_data = call_args[1]["data"] # data parameter from the update call # Verify team_member_budget is NOT in the update data assert ( "team_member_budget" not in update_data ), f"team_member_budget should not be in update data, but found: {update_data}" print( "✅ All test cases passed: team_member_budget is properly excluded from database update operations" ) def test_clean_team_member_fields(): """ Test that _clean_team_member_fields removes all team member fields from a dictionary. """ from litellm.proxy.management_endpoints.team_endpoints import ( TeamMemberBudgetHandler, ) data_dict = { "team_id": "test_team", "team_alias": "Test Team", "team_member_budget": 100.0, "team_member_budget_duration": "30d", "team_member_rpm_limit": 50, "team_member_tpm_limit": 1000, "other_field": "should_remain", } TeamMemberBudgetHandler._clean_team_member_fields(data_dict) assert "team_member_budget" not in data_dict assert "team_member_budget_duration" not in data_dict assert "team_member_rpm_limit" not in data_dict assert "team_member_tpm_limit" not in data_dict assert data_dict["team_id"] == "test_team" assert data_dict["team_alias"] == "Test Team" assert data_dict["other_field"] == "should_remain" def test_clean_team_member_fields_with_missing_fields(): """ Test that _clean_team_member_fields handles dictionaries without team member fields gracefully. """ from litellm.proxy.management_endpoints.team_endpoints import ( TeamMemberBudgetHandler, ) data_dict = { "team_id": "test_team", "team_alias": "Test Team", } TeamMemberBudgetHandler._clean_team_member_fields(data_dict) assert data_dict["team_id"] == "test_team" assert data_dict["team_alias"] == "Test Team" @pytest.mark.asyncio async def test_create_team_member_budget_table(): """ Test that create_team_member_budget_table creates a budget and adds it to metadata. """ from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import LitellmUserRoles, NewTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import ( TeamMemberBudgetHandler, ) mock_user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) data = NewTeamRequest( team_id="test_team_id", team_alias="Test Team", budget_duration="1mo", ) new_team_data_json = { "team_id": "test_team_id", "team_alias": "Test Team", "team_member_budget": 100.0, "team_member_budget_duration": "30d", "team_member_rpm_limit": 50, "team_member_tpm_limit": 1000, } mock_budget_response = MagicMock() mock_budget_response.budget_id = "budget_123" with patch( "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget", new_callable=AsyncMock ) as mock_new_budget: mock_new_budget.return_value = mock_budget_response result = await TeamMemberBudgetHandler.create_team_member_budget_table( data=data, new_team_data_json=new_team_data_json, user_api_key_dict=mock_user_api_key_dict, team_member_budget=100.0, team_member_rpm_limit=50, team_member_tpm_limit=1000, team_member_budget_duration="30d", ) assert mock_new_budget.called call_args = mock_new_budget.call_args budget_request = call_args[1]["budget_obj"] assert budget_request.max_budget == 100.0 assert budget_request.rpm_limit == 50 assert budget_request.tpm_limit == 1000 assert budget_request.budget_duration == "30d" assert budget_request.budget_id is not None assert "team-" in budget_request.budget_id assert "team_member_budget_id" in result["metadata"] assert result["metadata"]["team_member_budget_id"] == "budget_123" assert "team_member_budget" not in result assert "team_member_budget_duration" not in result assert "team_member_rpm_limit" not in result assert "team_member_tpm_limit" not in result @pytest.mark.asyncio async def test_create_team_member_budget_table_without_team_alias(): """ Test that create_team_member_budget_table generates budget_id correctly when team_alias is None. """ from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import LitellmUserRoles, NewTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import ( TeamMemberBudgetHandler, ) mock_user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) data = NewTeamRequest(team_id="test_team_id") new_team_data_json = { "team_id": "test_team_id", "team_member_budget": 100.0, } mock_budget_response = MagicMock() mock_budget_response.budget_id = "budget_123" with patch( "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget", new_callable=AsyncMock ) as mock_new_budget: mock_new_budget.return_value = mock_budget_response result = await TeamMemberBudgetHandler.create_team_member_budget_table( data=data, new_team_data_json=new_team_data_json, user_api_key_dict=mock_user_api_key_dict, team_member_budget=100.0, ) assert mock_new_budget.called call_args = mock_new_budget.call_args budget_request = call_args[1]["budget_obj"] assert budget_request.budget_id is not None assert budget_request.budget_id.startswith("team-budget-") @pytest.mark.asyncio async def test_upsert_team_member_budget_table_existing_budget(): """ Test that upsert_team_member_budget_table updates an existing budget when team_member_budget_id exists. """ from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import ( TeamMemberBudgetHandler, ) mock_user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) team_table = MagicMock(spec=LiteLLM_TeamTable) team_table.metadata = {"team_member_budget_id": "existing_budget_123"} updated_kv = { "team_id": "test_team_id", "team_member_budget": 200.0, "team_member_budget_duration": "60d", "team_member_rpm_limit": 100, } mock_budget_response = MagicMock() mock_budget_response.budget_id = "existing_budget_123" with patch( "litellm.proxy.management_endpoints.budget_management_endpoints.update_budget", new_callable=AsyncMock ) as mock_update_budget: mock_update_budget.return_value = mock_budget_response result = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=team_table, user_api_key_dict=mock_user_api_key_dict, updated_kv=updated_kv, team_member_budget=200.0, team_member_budget_duration="60d", team_member_rpm_limit=100, ) assert mock_update_budget.called call_args = mock_update_budget.call_args budget_request = call_args[1]["budget_obj"] assert budget_request.budget_id == "existing_budget_123" assert budget_request.max_budget == 200.0 assert budget_request.budget_duration == "60d" assert budget_request.rpm_limit == 100 assert "team_member_budget_id" in result["metadata"] assert result["metadata"]["team_member_budget_id"] == "existing_budget_123" assert "team_member_budget" not in result assert "team_member_budget_duration" not in result assert "team_member_rpm_limit" not in result @pytest.mark.asyncio async def test_upsert_team_member_budget_table_no_existing_budget(): """ Test that upsert_team_member_budget_table creates a new budget when team_member_budget_id does not exist. """ from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import ( TeamMemberBudgetHandler, ) mock_user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) team_table = MagicMock(spec=LiteLLM_TeamTable) team_table.metadata = {} team_table.team_alias = "Test Team" team_table.budget_duration = None updated_kv = { "team_id": "test_team_id", "team_member_budget": 150.0, "team_member_budget_duration": "45d", } mock_budget_response = MagicMock() mock_budget_response.budget_id = "new_budget_456" with patch( "litellm.proxy.management_endpoints.budget_management_endpoints.new_budget", new_callable=AsyncMock ) as mock_new_budget: mock_new_budget.return_value = mock_budget_response result = await TeamMemberBudgetHandler.upsert_team_member_budget_table( team_table=team_table, user_api_key_dict=mock_user_api_key_dict, updated_kv=updated_kv, team_member_budget=150.0, team_member_budget_duration="45d", ) assert mock_new_budget.called assert "team_member_budget_id" in result["metadata"] assert result["metadata"]["team_member_budget_id"] == "new_budget_456" assert "team_member_budget" not in result assert "team_member_budget_duration" not in result @pytest.mark.asyncio async def test_update_team_with_team_member_budget_duration(): """ Test that team/update endpoint properly handles team_member_budget_duration. """ from unittest.mock import AsyncMock, MagicMock, Mock, patch from fastapi import Request from litellm.proxy._types import LitellmUserRoles, UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team mock_request = Mock(spec=Request) mock_user_api_key_dict = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id" ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client, patch( "litellm.proxy.proxy_server.llm_router" ) as mock_llm_router, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.proxy_logging_obj" ) as mock_logging, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.auth.auth_checks._cache_team_object" ) as mock_cache_team, patch( "litellm.proxy.management_endpoints.team_endpoints.TeamMemberBudgetHandler.upsert_team_member_budget_table" ) as mock_upsert_budget: mock_existing_team = MagicMock() mock_existing_team.model_dump.return_value = { "team_id": "test_team_id", "team_alias": "test_team", "metadata": {"team_member_budget_id": "budget_123"}, } mock_existing_team.metadata = {"team_member_budget_id": "budget_123"} mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) mock_updated_team = MagicMock() mock_updated_team.team_id = "test_team_id" mock_updated_team.model_dump.return_value = {"team_id": "test_team_id"} mock_prisma_client.db.litellm_teamtable.update = AsyncMock( return_value=mock_updated_team ) mock_prisma_client.jsonify_team_object = MagicMock( side_effect=lambda db_data: db_data ) 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_member_budget_duration=None ): result_kv = updated_kv.copy() result_kv.pop("team_member_budget", None) result_kv.pop("team_member_budget_duration", None) return result_kv mock_upsert_budget.side_effect = mock_upsert_side_effect update_request = UpdateTeamRequest( team_id="test_team_id", team_alias="updated_alias", team_member_budget=100.0, team_member_budget_duration="30d", ) result = await update_team( data=update_request, http_request=mock_request, user_api_key_dict=mock_user_api_key_dict, ) assert mock_upsert_budget.called call_args = mock_upsert_budget.call_args assert call_args[1]["team_member_budget"] == 100.0 assert call_args[1]["team_member_budget_duration"] == "30d" assert mock_prisma_client.db.litellm_teamtable.update.called update_call_args = mock_prisma_client.db.litellm_teamtable.update.call_args update_data = update_call_args[1]["data"] assert "team_member_budget" not in update_data assert "team_member_budget_duration" not in update_data @pytest.mark.asyncio async def test_bulk_team_member_add_success(): """ Test bulk_team_member_add with successful addition of multiple members """ from litellm.proxy._types import ( LiteLLM_TeamMembership, LiteLLM_UserTable, TeamAddMemberResponse, ) from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add # Create test data test_members = [ Member(user_email="user1@example.com", role="user"), Member(user_email="user2@example.com", role="admin"), ] bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", members=test_members, max_budget_in_team=100.0, ) # Mock successful team_member_add response using MagicMock for simplicity mock_user_1 = MagicMock(spec=LiteLLM_UserTable) mock_user_1.user_id = "user-1" mock_user_1.user_email = "user1@example.com" mock_user_1.model_dump.return_value = { "user_id": "user-1", "user_email": "user1@example.com", } mock_user_2 = MagicMock(spec=LiteLLM_UserTable) mock_user_2.user_id = "user-2" mock_user_2.user_email = "user2@example.com" mock_user_2.model_dump.return_value = { "user_id": "user-2", "user_email": "user2@example.com", } mock_updated_users = [mock_user_1, mock_user_2] mock_membership_1 = MagicMock(spec=LiteLLM_TeamMembership) mock_membership_1.user_id = "user-1" mock_membership_1.team_id = "test-team-123" mock_membership_1.model_dump.return_value = { "user_id": "user-1", "team_id": "test-team-123", } mock_membership_2 = MagicMock(spec=LiteLLM_TeamMembership) mock_membership_2.user_id = "user-2" mock_membership_2.team_id = "test-team-123" mock_membership_2.model_dump.return_value = { "user_id": "user-2", "team_id": "test-team-123", } mock_updated_memberships = [mock_membership_1, mock_membership_2] # Create a mock response that has model_dump method mock_team_response = MagicMock() mock_team_response.team_id = "test-team-123" mock_team_response.team_alias = "Test Team" mock_team_response.updated_users = mock_updated_users mock_team_response.updated_team_memberships = mock_updated_memberships mock_team_response.model_dump.return_value = { "team_id": "test-team-123", "team_alias": "Test Team", "updated_users": [u.model_dump() for u in mock_updated_users], "updated_team_memberships": [m.model_dump() for m in mock_updated_memberships], } with patch( "litellm.proxy.management_endpoints.team_endpoints.team_member_add", 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( data=bulk_request, user_api_key_dict=mock_auth, ) # Verify the result structure assert isinstance(result, BulkTeamMemberAddResponse) assert result.team_id == "test-team-123" assert result.total_requested == 2 assert result.successful_additions == 2 assert result.failed_additions == 0 assert len(result.results) == 2 # Verify individual results for i, member_result in enumerate(result.results): assert isinstance(member_result, TeamMemberAddResult) assert member_result.success is True assert member_result.error is None assert member_result.user_email == test_members[i].user_email # Verify team_member_add was called with correct data mock_team_member_add.assert_called_once() call_args = mock_team_member_add.call_args[1]["data"] assert call_args.team_id == "test-team-123" assert call_args.member == test_members assert call_args.max_budget_in_team == 100.0 @pytest.mark.asyncio async def test_bulk_team_member_add_no_members_error(): """ Test bulk_team_member_add raises error when no members provided """ from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", members=[], # Empty list ) mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc_info: await bulk_team_member_add( data=bulk_request, user_api_key_dict=mock_auth, ) assert exc_info.value.status_code == 400 assert "At least one member is required" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_bulk_team_member_add_batch_size_limit(): """ Test bulk_team_member_add enforces maximum batch size limit """ from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add # Create more than 500 members (the max batch size) large_member_list = [ Member(user_email=f"user{i}@example.com", role="user") for i in range(501) ] bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", members=large_member_list, ) mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc_info: await bulk_team_member_add( data=bulk_request, user_api_key_dict=mock_auth, ) assert exc_info.value.status_code == 400 assert "Maximum 500 members can be added at once" in str(exc_info.value.detail) @pytest.mark.asyncio 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.management_endpoints.team_endpoints import bulk_team_member_add bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", all_users=True, max_budget_in_team=50.0, ) # Mock database users mock_db_users = [ MagicMock(user_id="user-1", user_email="user1@example.com"), MagicMock(user_id="user-2", user_email="user2@example.com"), ] mock_team_response = TeamAddMemberResponse( team_id="test-team-123", team_alias="Test Team", updated_users=[], updated_team_memberships=[], ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.management_endpoints.team_endpoints.team_member_add", 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 ) mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) result = await bulk_team_member_add( data=bulk_request, user_api_key_dict=mock_auth, ) # Verify that find_many was called to get all users mock_prisma.db.litellm_usertable.find_many.assert_called_once_with( order={"created_at": "desc"} ) # Verify team_member_add was called with users from database mock_team_member_add.assert_called_once() call_args = mock_team_member_add.call_args[1]["data"] assert call_args.team_id == "test-team-123" assert len(call_args.member) == 2 # Should have 2 members from mock_db_users assert call_args.max_budget_in_team == 50.0 @pytest.mark.asyncio async def test_bulk_team_member_add_failure_scenario(): """ Test bulk_team_member_add handles failures gracefully """ from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add test_members = [ Member(user_email="user1@example.com", role="user"), Member(user_email="user2@example.com", role="admin"), ] bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", members=test_members, ) with patch( "litellm.proxy.management_endpoints.team_endpoints.team_member_add", 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( data=bulk_request, user_api_key_dict=mock_auth, ) # Verify failure response structure assert isinstance(result, BulkTeamMemberAddResponse) assert result.team_id == "test-team-123" assert result.total_requested == 2 assert result.successful_additions == 0 assert result.failed_additions == 2 assert result.updated_team is None # Verify all members marked as failed assert len(result.results) == 2 for member_result in result.results: assert member_result.success is False assert member_result.error == "Database connection failed" @pytest.mark.asyncio async def test_bulk_team_member_add_no_db_connection(): """ Test bulk_team_member_add handles missing database connection """ from litellm.proxy.management_endpoints.team_endpoints import bulk_team_member_add bulk_request = BulkTeamMemberAddRequest( team_id="test-team-123", members=[Member(user_email="user1@example.com", role="user")], ) with patch("litellm.proxy.proxy_server.prisma_client", None): mock_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) with pytest.raises(HTTPException) as exc_info: await bulk_team_member_add( data=bulk_request, user_api_key_dict=mock_auth, ) assert exc_info.value.status_code == 500 assert "DB not connected" in str(exc_info.value.detail) @pytest.mark.asyncio 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 fastapi import HTTPException, Request from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 # Mock request mock_request = Mock(spec=Request) # Test Case 1: Non-admin user trying to query all teams (user_id=None) mock_user_api_key_dict_non_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non_admin_user_123", ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client # Should raise HTTPException with 401 status with pytest.raises(HTTPException) as exc_info: await list_team_v2( http_request=mock_request, user_id=None, # Non-admin trying to query all teams user_api_key_dict=mock_user_api_key_dict_non_admin, status=None, ) assert exc_info.value.status_code == 401 assert "Only admin users can query all teams/other teams" in str( exc_info.value.detail ) assert LitellmUserRoles.INTERNAL_USER.value in str(exc_info.value.detail) @pytest.mark.asyncio 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 fastapi import HTTPException, Request from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 # Mock request mock_request = Mock(spec=Request) # Test Case 2: Non-admin user trying to query another user's teams mock_user_api_key_dict_non_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non_admin_user_123", ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: mock_prisma_client.return_value = MagicMock() # Mock non-None prisma client # Should raise HTTPException with 401 status with pytest.raises(HTTPException) as exc_info: await list_team_v2( http_request=mock_request, user_id="other_user_456", # Non-admin trying to query other user's teams user_api_key_dict=mock_user_api_key_dict_non_admin, status=None, ) assert exc_info.value.status_code == 401 assert "Only admin users can query all teams/other teams" in str( exc_info.value.detail ) @pytest.mark.asyncio async def test_list_team_v2_security_check_non_admin_user_own_teams(): """ Test that list_team_v2 allows non-admin users to query their own teams. """ from unittest.mock import AsyncMock, Mock, patch from fastapi import Request from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 # Mock request mock_request = Mock(spec=Request) # Test Case 3: Non-admin user querying their own teams (should be allowed) mock_user_api_key_dict_non_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non_admin_user_123", ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: # 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 = { "user_id": "non_admin_user_123", "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"}), Mock(model_dump=lambda: {"team_id": "team_2", "team_alias": "Team 2"}), ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) # Should NOT raise an exception result = await list_team_v2( http_request=mock_request, user_id="non_admin_user_123", # Non-admin querying their own teams user_api_key_dict=mock_user_api_key_dict_non_admin, team_id=None, page=1, page_size=10, status=None, ) # Should return results without error assert "teams" in result assert "total" in result assert result["total"] == 2 @pytest.mark.asyncio async def test_list_team_v2_security_check_admin_user(): """ Test that list_team_v2 allows admin users to query any teams. """ from unittest.mock import AsyncMock, Mock, patch from fastapi import Request from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 # Mock request mock_request = Mock(spec=Request) # Test Case 4: Admin user querying all teams (should be allowed) mock_user_api_key_dict_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user_123", ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: # 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"}), Mock(model_dump=lambda: {"team_id": "team_2", "team_alias": "Team 2"}), ] mock_db.litellm_teamtable.find_many = AsyncMock(return_value=mock_teams) mock_db.litellm_teamtable.count = AsyncMock(return_value=2) # Should NOT raise an exception result = await list_team_v2( http_request=mock_request, user_id=None, # Admin querying all teams user_api_key_dict=mock_user_api_key_dict_admin, page=1, page_size=10, status=None, ) # Should return results without error assert "teams" in result assert "total" in result assert result["total"] == 2 @pytest.mark.asyncio async def test_list_team_v2_with_status_deleted(): """ Test that status="deleted" parameter correctly queries the deleted teams table. """ from unittest.mock import AsyncMock, Mock, patch from fastapi import Request from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 # Mock request mock_request = Mock(spec=Request) # Mock admin user mock_user_api_key_dict_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user_123", ) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma_client: # Mock prisma client and database operations mock_db = Mock() mock_prisma_client.db = mock_db # Mock deleted teams mock_deleted_team1 = Mock(model_dump=lambda: {"team_id": "team_1", "team_alias": "Deleted Team 1"}) mock_deleted_team2 = Mock(model_dump=lambda: {"team_id": "team_2", "team_alias": "Deleted Team 2"}) # Mock deleted teams table (should be called) mock_db.litellm_deletedteamtable.find_many = AsyncMock(return_value=[mock_deleted_team1, mock_deleted_team2]) mock_db.litellm_deletedteamtable.count = AsyncMock(return_value=2) # Mock regular teams table (should NOT be called) mock_db.litellm_teamtable.find_many = AsyncMock(return_value=[]) mock_db.litellm_teamtable.count = AsyncMock(return_value=0) # Should NOT raise an exception result = await list_team_v2( http_request=mock_request, user_id=None, # Admin querying all teams user_api_key_dict=mock_user_api_key_dict_admin, page=1, page_size=10, status="deleted", # Test the status parameter ) # Verify that deleted table was queried mock_db.litellm_deletedteamtable.find_many.assert_called_once() mock_db.litellm_deletedteamtable.count.assert_called_once() # Verify that regular table was NOT queried mock_db.litellm_teamtable.find_many.assert_not_called() mock_db.litellm_teamtable.count.assert_not_called() # Should return results without error assert "teams" in result assert "total" in result assert result["total"] == 2 assert len(result["teams"]) == 2 @pytest.mark.asyncio async def test_list_team_v2_with_invalid_status(): """ Test that invalid status parameter raises HTTPException. """ from unittest.mock import Mock, patch from fastapi import HTTPException, Request from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import list_team_v2 # Mock request mock_request = Mock(spec=Request) # Mock admin user mock_user_api_key_dict_admin = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user_123", ) mock_prisma_client = Mock() # Mock prisma_client to be non-None with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): # Should raise HTTPException for invalid status with pytest.raises(HTTPException) as exc_info: await list_team_v2( http_request=mock_request, user_id=None, user_api_key_dict=mock_user_api_key_dict_admin, page=1, page_size=10, status="invalid_status", # Invalid status value ) assert exc_info.value.status_code == 400 assert "Invalid status value" in str(exc_info.value.detail) assert "deleted" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_team_member_delete_cleans_membership(mock_db_client, mock_admin_auth): """ Verify that /team/member_delete removes the corresponding LiteLLM_TeamMembership row so the same user can be re-added without unique constraint issues. """ from litellm.proxy._types import TeamMemberDeleteRequest from litellm.proxy.management_endpoints.team_endpoints import team_member_delete test_team_id = "team-del-123" test_user_id = "user@example.com" # Mock Team row with the user as a member mock_team_row = MagicMock() mock_team_row.model_dump.return_value = { "team_id": test_team_id, "members_with_roles": [ {"user_id": test_user_id, "user_email": None, "role": "user"} ], "team_member_permissions": [], "metadata": {}, "models": [], "spend": 0.0, } # 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.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.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()) # Verification token deletion should be called mock_db_client.db.litellm_verificationtoken = MagicMock() mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) # Execute await team_member_delete( data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id), user_api_key_dict=mock_admin_auth, ) # Assert membership cleanup executed 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_team_member_delete_cleans_verification_tokens(mock_db_client, mock_admin_auth): from litellm.proxy._types import TeamMemberDeleteRequest from litellm.proxy.management_endpoints.team_endpoints import team_member_delete test_team_id = "team-del-tokens-123" test_user_id = "user-tokens@example.com" mock_team_row = MagicMock() mock_team_row.model_dump.return_value = { "team_id": test_team_id, "members_with_roles": [ {"user_id": test_user_id, "user_email": None, "role": "user"} ], "team_member_permissions": [], "metadata": {}, "models": [], "spend": 0.0, } 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) 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.update = AsyncMock(return_value=MagicMock()) mock_db_client.db.litellm_teammembership = MagicMock() mock_db_client.db.litellm_teammembership.delete_many = AsyncMock(return_value=MagicMock()) mock_db_client.db.litellm_verificationtoken = MagicMock() mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) mock_db_client.db.litellm_verificationtoken.delete_many = AsyncMock(return_value=MagicMock()) await team_member_delete( data=TeamMemberDeleteRequest(team_id=test_team_id, user_id=test_user_id), user_api_key_dict=mock_admin_auth, ) mock_db_client.db.litellm_verificationtoken.delete_many.assert_awaited_once_with( where={ "user_id": {"in": [test_user_id]}, "team_id": test_team_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. """ from fastapi import Request from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import new_team # Create non-admin user with user_max_budget set to 100.0 non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-user-123", user_max_budget=100.0, ) # Create team request with max_budget (200.0) exceeding user's limit (100.0) team_request = NewTeamRequest( team_alias="high-budget-team", max_budget=200.0, # Exceeds user's user_max_budget ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: # Setup basic mocks 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( data=team_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details # ProxyException stores status_code in 'code' attribute 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 LitellmUserRoles.INTERNAL_USER.value in str(exc_info.value.message) @pytest.mark.asyncio 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 from litellm.proxy._types import NewTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import new_team # Create non-admin user with user_max_budget set to 100.0 non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-user-456", user_max_budget=100.0, models=[], # Empty models list to bypass model validation ) # Create team request with max_budget (50.0) within user's limit (100.0) team_request = NewTeamRequest( team_alias="within-budget-team", max_budget=50.0, # Within user's user_max_budget ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), 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" mock_created_team.team_alias = "within-budget-team" mock_created_team.max_budget = 50.0 mock_created_team.members_with_roles = [] mock_created_team.metadata = None mock_created_team.model_dump.return_value = { "team_id": "team-within-budget-789", "team_alias": "within-budget-team", "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 model table mock_prisma.db.litellm_modeltable = MagicMock() 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_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 = { "team_id": "team-within-budget-789", "user_id": "non-admin-user-456", "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) # Should NOT raise an exception result = await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify the team was created successfully 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_new_team_org_scoped_budget_bypasses_user_limit(): """ Test that /team/new with organization_id does NOT validate budget against user's personal max_budget. This is the bug fix for: When an org admin creates an org-scoped team, the team's budget should be validated against the organization's limits, not the user's personal limits. Scenario: - Organization has max_budget=$100 - User (org admin) has personal max_budget=$3 - Team is created with organization_id and max_budget=$50 - Expected: Should succeed (within org's $100 limit) - Bug behavior: Would fail with "max budget higher than user max. User max budget=3.0" """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_OrganizationTable, LiteLLM_UserTable, NewTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create non-admin user with very restrictive personal budget ($3) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-user-123", user_max_budget=3.0, # Restrictive personal budget models=[], # Empty models list to bypass model validation ) # Create team request with budget ($50) that's within org's limit but exceeds user's personal limit team_request = NewTeamRequest( team_alias="org-scoped-team", max_budget=50.0, # Within org's $100 limit, but exceeds user's $3 limit organization_id="test-org-123", # This makes it an org-scoped team ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object" ) as mock_get_org: # 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 organization with $100 budget mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-123" mock_org.max_budget = 100.0 mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] mock_org.litellm_budget_table = None # No budget table for this test mock_get_org.return_value = mock_org # Mock user cache to return user with restrictive personal budget mock_user_obj = LiteLLM_UserTable( user_id="org-admin-user-123", max_budget=3.0, # Restrictive personal budget ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Mock team creation mock_created_team = MagicMock() mock_created_team.team_id = "team-org-scoped-789" mock_created_team.team_alias = "org-scoped-team" mock_created_team.max_budget = 50.0 mock_created_team.organization_id = "test-org-123" mock_created_team.members_with_roles = [] mock_created_team.metadata = None mock_created_team.model_dump.return_value = { "team_id": "team-org-scoped-789", "team_alias": "org-scoped-team", "max_budget": 50.0, "organization_id": "test-org-123", "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 model table mock_prisma.db.litellm_modeltable = MagicMock() mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations mock_user = MagicMock() mock_user.user_id = "org-admin-user-123" mock_user.model_dump.return_value = {"user_id": "org-admin-user-123", "teams": ["team-org-scoped-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 = { "team_id": "team-org-scoped-789", "user_id": "org-admin-user-123", "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) # Should NOT raise an exception - the fix should bypass user budget validation for org-scoped teams result = await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify the team was created successfully with the higher budget assert result is not None assert result["team_id"] == "team-org-scoped-789" assert result["max_budget"] == 50.0 assert result["organization_id"] == "test-org-123" @pytest.mark.asyncio async def test_new_team_org_scoped_models_bypasses_user_limit(): """ Test that /team/new with organization_id does NOT validate models against user's personal models. This is the bug fix for: When an org admin creates an org-scoped team, the team's models should be validated against the organization's models, not the user's personal models. Scenario: - Organization has models=['gpt-4', 'gpt-3.5-turbo', 'claude-3-opus'] - User (org admin) has personal models=['no-default-models'] - Team is created with organization_id and models=['gpt-4'] - Expected: Should succeed (within org's allowed models) - Bug behavior: Would fail with "Model not in allowed user models" """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_OrganizationTable, LiteLLM_UserTable, NewTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create non-admin user with restrictive personal models org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-user-456", user_max_budget=None, # No budget restriction for this test models=["no-default-models"], # Restrictive personal models ) # Create team request with models that are within org's allowed models but not user's team_request = NewTeamRequest( team_alias="org-scoped-models-team", models=["gpt-4"], # Within org's allowed models, but not in user's personal models organization_id="test-org-456", # This makes it an org-scoped team ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object" ) as mock_get_org: # 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 organization with allowed models mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-456" mock_org.max_budget = 100.0 mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] mock_org.litellm_budget_table = None mock_get_org.return_value = mock_org # Mock user cache mock_user_obj = LiteLLM_UserTable( user_id="org-admin-user-456", max_budget=None, ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Mock team creation mock_created_team = MagicMock() mock_created_team.team_id = "team-org-scoped-models-789" mock_created_team.team_alias = "org-scoped-models-team" mock_created_team.max_budget = None mock_created_team.organization_id = "test-org-456" mock_created_team.models = ["gpt-4"] mock_created_team.members_with_roles = [] mock_created_team.metadata = None mock_created_team.model_dump.return_value = { "team_id": "team-org-scoped-models-789", "team_alias": "org-scoped-models-team", "max_budget": None, "organization_id": "test-org-456", "models": ["gpt-4"], "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 model table mock_prisma.db.litellm_modeltable = MagicMock() mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations mock_user = MagicMock() mock_user.user_id = "org-admin-user-456" mock_user.model_dump.return_value = {"user_id": "org-admin-user-456", "teams": ["team-org-scoped-models-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 = { "team_id": "team-org-scoped-models-789", "user_id": "org-admin-user-456", "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) # Should NOT raise an exception - the fix should bypass user model validation for org-scoped teams result = await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify the team was created successfully with the org's models assert result is not None assert result["team_id"] == "team-org-scoped-models-789" assert result["models"] == ["gpt-4"] assert result["organization_id"] == "test-org-456" @pytest.mark.asyncio async def test_new_team_standalone_validates_against_user_models(): """ Test that /team/new WITHOUT organization_id still validates models against user's personal models. This ensures that standalone teams (not org-scoped) still use user-level validation. Scenario: - User has personal models=['no-default-models'] - Team is created WITHOUT organization_id and models=['gpt-4'] - Expected: Should fail with "Model not in allowed user models" """ from fastapi import Request from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import new_team # Create non-admin user with restrictive personal models non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-user-789", user_max_budget=None, models=["no-default-models"], # Restrictive personal models ) # Create standalone team request (no organization_id) with models not in user's list team_request = NewTeamRequest( team_alias="standalone-team", models=["gpt-4"], # Not in user's allowed models # Note: No organization_id - this is a standalone team ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: # Setup basic mocks 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) # Should raise ProxyException because gpt-4 is not in user's allowed models with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "Model not in allowed user models" in str(exc_info.value.message) assert "no-default-models" in str(exc_info.value.message) @pytest.mark.asyncio async def test_new_team_standalone_validates_against_user_budget(): """ Test that /team/new WITHOUT organization_id still validates budget against user's personal max_budget. This ensures that standalone teams (not org-scoped) still use user-level validation. This is essentially the same as test_new_team_max_budget_exceeds_user_max_budget but explicitly showing the contrast with org-scoped teams. Scenario: - User has personal max_budget=$3 - Team is created WITHOUT organization_id and max_budget=$50 - Expected: Should fail with "max budget higher than user max" """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_UserTable, NewTeamRequest, ProxyException, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create non-admin user with restrictive personal budget non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-user-budget-789", user_max_budget=100.0, # This is for key auth, actual budget is from user object models=[], # Empty models list to bypass model validation ) # Create standalone team request (no organization_id) with budget exceeding user's limit team_request = NewTeamRequest( team_alias="standalone-budget-team", max_budget=50.0, # Exceeds user's personal budget # Note: No organization_id - this is a standalone team ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: # Setup basic mocks 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 user with restrictive personal budget ($3) mock_user_obj = LiteLLM_UserTable( user_id="non-admin-user-budget-789", max_budget=3.0, # Restrictive personal budget ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Should raise ProxyException because budget exceeds user's max_budget with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "max budget higher than user max" in str(exc_info.value.message) assert "3.0" in str(exc_info.value.message) # User's max_budget should be mentioned @pytest.mark.asyncio async def test_new_team_org_scoped_budget_exceeds_org_limit(): """ Test that /team/new with organization_id fails when team budget exceeds organization's max_budget. Scenario: - Organization has max_budget=$100 - Team is created with organization_id and max_budget=$150 - Expected: Should fail with error about exceeding org budget """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, NewTeamRequest, ProxyException, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create user (org admin) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-user-budget-test", models=[], ) # Create team request with budget ($150) that exceeds org's limit ($100) team_request = NewTeamRequest( team_alias="org-team-exceeds-budget", max_budget=150.0, # Exceeds org's $100 limit organization_id="test-org-budget-limit", ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object" ) as mock_get_org: # Setup mocks 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 organization with $100 budget limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.max_budget = 100.0 mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-budget-limit" mock_org.models = ["gpt-4", "gpt-3.5-turbo"] mock_org.litellm_budget_table = mock_budget_table mock_get_org.return_value = mock_org # Should raise ProxyException because team budget exceeds org budget with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "exceeds organization" in str(exc_info.value.message).lower() or "organization" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_new_team_org_scoped_models_not_in_org_models(): """ Test that /team/new with organization_id fails when team models are not in organization's allowed models. Scenario: - Organization has models=['gpt-4', 'gpt-3.5-turbo'] - Team is created with organization_id and models=['claude-3-opus'] - Expected: Should fail with error about model not in org's allowed models """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, NewTeamRequest, ProxyException, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create user (org admin) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-user-models-test", models=[], ) # Create team request with model not in org's allowed list team_request = NewTeamRequest( team_alias="org-team-invalid-model", models=["claude-3-opus"], # Not in org's allowed models organization_id="test-org-models-limit", ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object" ) as mock_get_org: # Setup mocks 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 organization with specific allowed models (not including claude-3-opus) mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-models-limit" mock_org.models = ["gpt-4", "gpt-3.5-turbo"] # claude-3-opus is NOT allowed mock_org.litellm_budget_table = None mock_get_org.return_value = mock_org # Should raise ProxyException because claude-3-opus is not in org's allowed models with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "claude-3-opus" in str(exc_info.value.message) or "organization" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_standalone_budget_exceeds_user_limit(): """ Test that /team/update for a standalone team fails when new budget exceeds user's max_budget. Scenario: - User has personal max_budget=$50 - Standalone team exists (no organization_id) - User tries to update team budget to $100 - Expected: Should fail with error about exceeding user budget """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_UserTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create non-admin user with restrictive personal budget non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-update-test", models=[], ) # Create update request with budget exceeding user's limit update_request = UpdateTeamRequest( team_id="standalone-team-123", max_budget=100.0, # Exceeds user's $50 limit ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: # Mock existing standalone team (no organization_id) mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-123" mock_existing_team.organization_id = None # Standalone team mock_existing_team.max_budget = 30.0 mock_existing_team.model_dump.return_value = { "team_id": "standalone-team-123", "organization_id": None, "max_budget": 30.0, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Mock user cache to return user with restrictive budget mock_user_obj = LiteLLM_UserTable( user_id="non-admin-update-test", max_budget=50.0, # User's budget limit ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Should raise ProxyException because new budget exceeds user's max_budget with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "budget" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_org_scoped_budget_exceeds_org_limit(): """ Test that /team/update for an org-scoped team fails when new budget exceeds organization's max_budget. Scenario: - Organization has max_budget=$100 - Org-scoped team exists - User tries to update team budget to $150 - Expected: Should fail with error about exceeding org budget """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user (org admin) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-test", models=[], ) # Create update request with budget exceeding org's limit update_request = UpdateTeamRequest( team_id="org-team-456", max_budget=150.0, # Exceeds org's $100 limit ) dummy_request = MagicMock(spec=Request) # Mock organization with $100 budget limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.max_budget = 100.0 mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ) as mock_get_org: # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-456" mock_existing_team.organization_id = "test-org-update" mock_existing_team.max_budget = 80.0 mock_existing_team.model_dump.return_value = { "team_id": "org-team-456", "organization_id": "test-org-update", "max_budget": 80.0, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because new budget exceeds org's max_budget with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "organization" in str(exc_info.value.message).lower() or "budget" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_standalone_models_exceeds_user_limit(): """ Test that /team/update for a standalone team fails when models are not in user's allowed models. Scenario: - User has personal models=['gpt-3.5-turbo'] - Standalone team exists (no organization_id) - User tries to update team models to ['gpt-4'] (not in user's allowed models) - Expected: Should fail with error about model not in user's allowed models """ from fastapi import Request from litellm.proxy._types import ProxyException, UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team # Create non-admin user with restrictive personal models non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="non-admin-update-models-test", models=["gpt-3.5-turbo"], # Restrictive model list ) # Create update request with model not in user's allowed list update_request = UpdateTeamRequest( team_id="standalone-team-models-123", models=["gpt-4"], # Not in user's allowed models ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: # Mock existing standalone team (no organization_id) mock_existing_team = MagicMock() mock_existing_team.team_id = "standalone-team-models-123" mock_existing_team.organization_id = None # Standalone team mock_existing_team.models = ["gpt-3.5-turbo"] mock_existing_team.model_dump.return_value = { "team_id": "standalone-team-models-123", "organization_id": None, "models": ["gpt-3.5-turbo"], } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because model not in user's allowed models with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "model" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_org_scoped_budget_bypasses_user_limit(): """ Test that /team/update for an org-scoped team does NOT validate budget against user's personal max_budget. Scenario: - Organization has max_budget=$100 - User (org admin) has personal max_budget=$3 - Org-scoped team exists with current budget=$30 - User tries to update team budget to $50 (within org limit, exceeds user limit) - Expected: Should succeed (validated against org, not user) """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, LiteLLM_UserTable, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user with very restrictive personal budget ($3) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-budget-test", models=[], ) # Create update request with budget within org limit but exceeding user limit update_request = UpdateTeamRequest( team_id="org-team-update-budget-123", max_budget=50.0, # Within org's $100 limit, exceeds user's $3 limit ) dummy_request = MagicMock(spec=Request) # Mock organization with $100 budget limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.max_budget = 100.0 mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update-budget" mock_org.models = ["gpt-4", "gpt-3.5-turbo"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ) as mock_get_org: # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-budget-123" mock_existing_team.organization_id = "test-org-update-budget" mock_existing_team.max_budget = 30.0 mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "org-team-update-budget-123", "organization_id": "test-org-update-budget", "max_budget": 30.0, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data # Mock user cache to return user with restrictive budget mock_user_obj = LiteLLM_UserTable( user_id="org-admin-update-budget-test", max_budget=3.0, # Restrictive personal budget ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object # Mock team update mock_updated_team = MagicMock() mock_updated_team.team_id = "org-team-update-budget-123" mock_updated_team.organization_id = "test-org-update-budget" mock_updated_team.max_budget = 50.0 mock_updated_team.litellm_model_table = None mock_updated_team.model_dump.return_value = { "team_id": "org-team-update-budget-123", "organization_id": "test-org-update-budget", "max_budget": 50.0, } mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise an exception - bypass user budget validation for org-scoped teams result = await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify the team was updated successfully with the higher budget assert result is not None assert result["data"].max_budget == 50.0 @pytest.mark.asyncio async def test_update_team_org_scoped_models_bypasses_user_limit(): """ Test that /team/update for an org-scoped team does NOT validate models against user's personal models. Scenario: - Organization has models=['gpt-4', 'gpt-3.5-turbo', 'claude-3-opus'] - User (org admin) has personal models=['no-default-models'] - Org-scoped team exists - User tries to update team models to ['gpt-4'] (in org's allowed, not in user's) - Expected: Should succeed (validated against org, not user) """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_OrganizationTable, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user with very restrictive personal models org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-models-test", models=["no-default-models"], # Restrictive model list ) # Create update request with models in org's allowed but not in user's update_request = UpdateTeamRequest( team_id="org-team-update-models-123", models=["gpt-4"], # In org's allowed, not in user's ) dummy_request = MagicMock(spec=Request) # Mock organization with generous model list mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update-models" mock_org.models = ["gpt-4", "gpt-3.5-turbo", "claude-3-opus"] mock_org.litellm_budget_table = None with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ) as mock_get_org: # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-models-123" mock_existing_team.organization_id = "test-org-update-models" mock_existing_team.models = ["gpt-3.5-turbo"] mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "org-team-update-models-123", "organization_id": "test-org-update-models", "models": ["gpt-3.5-turbo"], } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object # Mock team update mock_updated_team = MagicMock() mock_updated_team.team_id = "org-team-update-models-123" mock_updated_team.organization_id = "test-org-update-models" mock_updated_team.models = ["gpt-4"] mock_updated_team.litellm_model_table = None mock_updated_team.model_dump.return_value = { "team_id": "org-team-update-models-123", "organization_id": "test-org-update-models", "models": ["gpt-4"], } mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise an exception - bypass user models validation for org-scoped teams result = await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify the team was updated successfully with the new models assert result is not None assert result["data"].models == ["gpt-4"] @pytest.mark.asyncio async def test_update_team_org_scoped_models_not_in_org_models(): """ Test that /team/update for an org-scoped team fails when models are not in organization's allowed models. Scenario: - Organization has models=['gpt-4', 'gpt-3.5-turbo'] - Org-scoped team exists - User tries to update team models to ['claude-3-opus'] (not in org's allowed models) - Expected: Should fail with error about model not in org's allowed models """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_OrganizationTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user (org admin) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-models-fail-test", models=[], ) # Create update request with model not in org's allowed list update_request = UpdateTeamRequest( team_id="org-team-update-models-fail-123", models=["claude-3-opus"], # Not in org's allowed models ) dummy_request = MagicMock(spec=Request) # Mock organization with restricted model list (no claude-3-opus) mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update-models-fail" mock_org.models = ["gpt-4", "gpt-3.5-turbo"] # claude-3-opus is NOT allowed mock_org.litellm_budget_table = None with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ) as mock_get_org: # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-models-fail-123" mock_existing_team.organization_id = "test-org-update-models-fail" mock_existing_team.models = ["gpt-4"] mock_existing_team.model_dump.return_value = { "team_id": "org-team-update-models-fail-123", "organization_id": "test-org-update-models-fail", "models": ["gpt-4"], } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because claude-3-opus is not in org's allowed models with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "claude-3-opus" in str(exc_info.value.message) or "organization" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_org_scoped_models_with_all_proxy_models(): """ Test that /team/update for an org-scoped team succeeds when organization has 'all-proxy-models'. Scenario: - Organization has models=['all-proxy-models'] (catch-all for all models) - Org-scoped team exists - User tries to update team models to ['rerank-english-v3.0', 'text-embedding-3-small', 'gpt-4o-mini-test'] - Expected: Should succeed because 'all-proxy-models' allows all models """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_OrganizationTable, SpecialModelNames, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user (org admin) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-all-proxy-models-test", models=[], ) # Create update request with models that aren't explicitly in org's models list # but should be allowed because org has 'all-proxy-models' update_request = UpdateTeamRequest( team_id="org-team-all-proxy-models-123", models=["rerank-english-v3.0", "text-embedding-3-small", "gpt-4o-mini-test"], ) dummy_request = MagicMock(spec=Request) # Mock organization with 'all-proxy-models' (catch-all) mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-all-proxy-models" mock_org.models = [SpecialModelNames.all_proxy_models.value] # Allows all models mock_org.litellm_budget_table = None with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ) as mock_get_org: # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-all-proxy-models-123" mock_existing_team.organization_id = "test-org-all-proxy-models" mock_existing_team.models = ["gpt-4"] mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "org-team-all-proxy-models-123", "organization_id": "test-org-all-proxy-models", "models": ["gpt-4"], } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object # Mock team update mock_updated_team = MagicMock() mock_updated_team.team_id = "org-team-all-proxy-models-123" mock_updated_team.organization_id = "test-org-all-proxy-models" mock_updated_team.models = ["rerank-english-v3.0", "text-embedding-3-small", "gpt-4o-mini-test"] mock_updated_team.litellm_model_table = None mock_updated_team.model_dump.return_value = { "team_id": "org-team-all-proxy-models-123", "organization_id": "test-org-all-proxy-models", "models": ["rerank-english-v3.0", "text-embedding-3-small", "gpt-4o-mini-test"], } mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) # Should NOT raise an exception - 'all-proxy-models' allows all models result = await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify the team was updated successfully with the new models assert result is not None assert result["data"].models == ["rerank-english-v3.0", "text-embedding-3-small", "gpt-4o-mini-test"] @pytest.mark.asyncio async def test_update_team_tpm_limit_exceeds_user_limit(): """ Test that /team/update fails when TPM limit exceeds user's TPM limit. Scenario: - User has tpm_limit=1000 - User tries to update team with tpm_limit=5000 - Expected: Should fail with error about exceeding user TPM limit """ from fastapi import Request from litellm.proxy._types import ProxyException, UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team # Create non-admin user with TPM limit non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="tpm-limit-user", models=[], tpm_limit=1000, # User's TPM limit ) # Create update request with TPM exceeding user's limit update_request = UpdateTeamRequest( team_id="team-tpm-test-123", tpm_limit=5000, # Exceeds user's 1000 limit ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ): # Mock existing standalone team mock_existing_team = MagicMock() mock_existing_team.team_id = "team-tpm-test-123" mock_existing_team.organization_id = None mock_existing_team.tpm_limit = 500 mock_existing_team.model_dump.return_value = { "team_id": "team-tpm-test-123", "organization_id": None, "tpm_limit": 500, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because new TPM exceeds user's limit with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "tpm" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_rpm_limit_exceeds_user_limit(): """ Test that /team/update fails when RPM limit exceeds user's RPM limit. Scenario: - User has rpm_limit=100 - User tries to update team with rpm_limit=500 - Expected: Should fail with error about exceeding user RPM limit """ from fastapi import Request from litellm.proxy._types import ProxyException, UpdateTeamRequest, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import update_team # Create non-admin user with RPM limit non_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="rpm-limit-user", models=[], rpm_limit=100, # User's RPM limit ) # Create update request with RPM exceeding user's limit update_request = UpdateTeamRequest( team_id="team-rpm-test-123", rpm_limit=500, # Exceeds user's 100 limit ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ): # Mock existing standalone team mock_existing_team = MagicMock() mock_existing_team.team_id = "team-rpm-test-123" mock_existing_team.organization_id = None mock_existing_team.rpm_limit = 50 mock_existing_team.model_dump.return_value = { "team_id": "team-rpm-test-123", "organization_id": None, "rpm_limit": 50, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because new RPM exceeds user's limit with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=non_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "rpm" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_new_team_org_scoped_tpm_exceeds_org_limit(): """ Test that /team/new for an org-scoped team fails when TPM exceeds organization's TPM limit. Scenario: - Organization has tpm_limit=10000 - User tries to create org-scoped team with tpm_limit=20000 - Expected: Should fail with error about exceeding org TPM limit """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, NewTeamRequest, ProxyException, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create user (with restrictive personal TPM limit that should be bypassed) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-tpm-test", models=[], tpm_limit=1000, # User's personal limit (should be bypassed for org teams) ) # Create team request with TPM exceeding org's limit team_request = NewTeamRequest( team_alias="org-tpm-test-team", organization_id="test-org-tpm", tpm_limit=20000, # Exceeds org's 10000 limit ) dummy_request = MagicMock(spec=Request) # Mock organization with TPM limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.tpm_limit = 10000 # Org's TPM limit mock_budget_table.rpm_limit = None mock_budget_table.max_budget = None mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-tpm" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ): mock_license.is_team_count_over_limit.return_value = False mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) # Should raise ProxyException because TPM exceeds org limit with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "tpm" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_new_team_org_scoped_rpm_exceeds_org_limit(): """ Test that /team/new for an org-scoped team fails when RPM exceeds organization's RPM limit. Scenario: - Organization has rpm_limit=1000 - User tries to create org-scoped team with rpm_limit=2000 - Expected: Should fail with error about exceeding org RPM limit """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, NewTeamRequest, ProxyException, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create user (with restrictive personal RPM limit that should be bypassed) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-rpm-test", models=[], rpm_limit=100, # User's personal limit (should be bypassed for org teams) ) # Create team request with RPM exceeding org's limit team_request = NewTeamRequest( team_alias="org-rpm-test-team", organization_id="test-org-rpm", rpm_limit=2000, # Exceeds org's 1000 limit ) dummy_request = MagicMock(spec=Request) # Mock organization with RPM limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.tpm_limit = None mock_budget_table.rpm_limit = 1000 # Org's RPM limit mock_budget_table.max_budget = None mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-rpm" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ): mock_license.is_team_count_over_limit.return_value = False mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) # Should raise ProxyException because RPM exceeds org limit with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "rpm" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit(): """ Test that /team/new for an org-scoped team bypasses user's TPM/RPM limits. Scenario: - User has tpm_limit=1000, rpm_limit=100 - Organization has tpm_limit=50000, rpm_limit=5000 - User creates org-scoped team with tpm_limit=10000, rpm_limit=1000 - Expected: Should succeed (bypasses user limits, within org limits) """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, LiteLLM_TeamTable, NewTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import new_team # Create user with restrictive personal limits org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-bypass-test", models=[], tpm_limit=1000, # Restrictive user TPM limit rpm_limit=100, # Restrictive user RPM limit ) # Create team request exceeding user limits but within org limits team_request = NewTeamRequest( team_alias="org-bypass-test-team", organization_id="test-org-bypass", tpm_limit=10000, # Exceeds user's 1000 but within org's 50000 rpm_limit=1000, # Exceeds user's 100 but within org's 5000 ) dummy_request = MagicMock(spec=Request) # Mock organization with generous limits mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.tpm_limit = 50000 # Generous org TPM limit mock_budget_table.rpm_limit = 5000 # Generous org RPM limit mock_budget_table.max_budget = None mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-bypass" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ), patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ), patch( "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", new=AsyncMock() ): mock_license.is_team_count_over_limit.return_value = False mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0) mock_prisma.get_data = AsyncMock(return_value=None) # Mock team creation mock_created_team = MagicMock(spec=LiteLLM_TeamTable) mock_created_team.team_id = "new-bypass-team-id" mock_created_team.team_alias = "org-bypass-test-team" mock_created_team.tpm_limit = 10000 mock_created_team.rpm_limit = 1000 mock_created_team.metadata = None mock_created_team.members_with_roles = [] mock_created_team.model_dump.return_value = { "team_id": "new-bypass-team-id", "team_alias": "org-bypass-test-team", "tpm_limit": 10000, "rpm_limit": 1000, "metadata": None, "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.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # Should succeed - bypasses user limits since org-scoped result = await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify team was created assert result["team_id"] == "new-bypass-team-id" @pytest.mark.asyncio async def test_update_team_org_scoped_tpm_exceeds_org_limit(): """ Test that /team/update for an org-scoped team fails when TPM exceeds organization's TPM limit. Scenario: - Organization has tpm_limit=10000 - User tries to update org-scoped team with tpm_limit=20000 - Expected: Should fail with error about exceeding org TPM limit """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user (with restrictive personal TPM limit that should be bypassed) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-tpm-test", models=[], tpm_limit=1000, # User's personal limit (should be bypassed for org teams) ) # Create update request with TPM exceeding org's limit update_request = UpdateTeamRequest( team_id="org-team-update-tpm-123", tpm_limit=20000, # Exceeds org's 10000 limit ) dummy_request = MagicMock(spec=Request) # Mock organization with TPM limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.tpm_limit = 10000 # Org's TPM limit mock_budget_table.rpm_limit = None mock_budget_table.max_budget = None mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update-tpm" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-tpm-123" mock_existing_team.organization_id = "test-org-update-tpm" mock_existing_team.tpm_limit = 5000 mock_existing_team.model_dump.return_value = { "team_id": "org-team-update-tpm-123", "organization_id": "test-org-update-tpm", "tpm_limit": 5000, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because TPM exceeds org limit with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "tpm" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_org_scoped_rpm_exceeds_org_limit(): """ Test that /team/update for an org-scoped team fails when RPM exceeds organization's RPM limit. Scenario: - Organization has rpm_limit=1000 - User tries to update org-scoped team with rpm_limit=2000 - Expected: Should fail with error about exceeding org RPM limit """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user (with restrictive personal RPM limit that should be bypassed) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-rpm-test", models=[], rpm_limit=100, # User's personal limit (should be bypassed for org teams) ) # Create update request with RPM exceeding org's limit update_request = UpdateTeamRequest( team_id="org-team-update-rpm-123", rpm_limit=2000, # Exceeds org's 1000 limit ) dummy_request = MagicMock(spec=Request) # Mock organization with RPM limit mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.tpm_limit = None mock_budget_table.rpm_limit = 1000 # Org's RPM limit mock_budget_table.max_budget = None mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update-rpm" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-rpm-123" mock_existing_team.organization_id = "test-org-update-rpm" mock_existing_team.rpm_limit = 500 mock_existing_team.model_dump.return_value = { "team_id": "org-team-update-rpm-123", "organization_id": "test-org-update-rpm", "rpm_limit": 500, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Should raise ProxyException because RPM exceeds org limit with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify exception details assert exc_info.value.code == '400' assert "rpm" in str(exc_info.value.message).lower() @pytest.mark.asyncio async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit(): """ Test that /team/update for an org-scoped team bypasses user's TPM/RPM limits. Scenario: - User has tpm_limit=1000, rpm_limit=100 - Organization has tpm_limit=50000, rpm_limit=5000 - User updates org-scoped team with tpm_limit=10000, rpm_limit=1000 - Expected: Should succeed (bypasses user limits, within org limits) """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_BudgetTable, LiteLLM_OrganizationTable, LiteLLM_TeamTable, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user with restrictive personal limits org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-update-bypass-test", models=[], tpm_limit=1000, # Restrictive user TPM limit rpm_limit=100, # Restrictive user RPM limit ) # Create update request exceeding user limits but within org limits update_request = UpdateTeamRequest( team_id="org-team-update-bypass-123", tpm_limit=10000, # Exceeds user's 1000 but within org's 50000 rpm_limit=1000, # Exceeds user's 100 but within org's 5000 ) dummy_request = MagicMock(spec=Request) # Mock organization with generous limits mock_budget_table = MagicMock(spec=LiteLLM_BudgetTable) mock_budget_table.tpm_limit = 50000 # Generous org TPM limit mock_budget_table.rpm_limit = 5000 # Generous org RPM limit mock_budget_table.max_budget = None mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-update-bypass" mock_org.models = ["gpt-4"] mock_org.litellm_budget_table = mock_budget_table with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.proxy_logging_obj" ) as mock_logging, patch( "litellm.proxy.management_endpoints.team_endpoints.get_org_object", new=AsyncMock(return_value=mock_org) ): # Mock existing org-scoped team mock_existing_team = MagicMock() mock_existing_team.team_id = "org-team-update-bypass-123" mock_existing_team.organization_id = "test-org-update-bypass" mock_existing_team.tpm_limit = 5000 mock_existing_team.rpm_limit = 500 mock_existing_team.model_id = None mock_existing_team.model_dump.return_value = { "team_id": "org-team-update-bypass-123", "organization_id": "test-org-update-bypass", "tpm_limit": 5000, "rpm_limit": 500, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) mock_cache.async_set_cache = AsyncMock() # Mock team update mock_updated_team = MagicMock(spec=LiteLLM_TeamTable) mock_updated_team.team_id = "org-team-update-bypass-123" mock_updated_team.tpm_limit = 10000 mock_updated_team.rpm_limit = 1000 mock_updated_team.model_dump.return_value = { "team_id": "org-team-update-bypass-123", "tpm_limit": 10000, "rpm_limit": 1000, } mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # Should succeed - bypasses user limits since org-scoped result = await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify team was updated assert result["team_id"] == "org-team-update-bypass-123" @pytest.mark.asyncio async def test_update_team_guardrails_with_org_id(): """ Test that updating team guardrails works when team has an organization_id. The fix ensures 'teams' field is included when fetching organization data. """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_OrganizationTable, LiteLLM_TeamTable, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create user (org admin) org_admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.INTERNAL_USER, user_id="org-admin-guardrails-test", models=[], ) # Update request to add guardrails to team update_request = UpdateTeamRequest( team_id="team-guardrails-123", guardrails=["aporia-pre-call", "aporia-post-call"], organization_id="test-org-guardrails", # Changing org triggers fetch_and_validate_organization ) dummy_request = MagicMock(spec=Request) # Mock organization with all required fields including teams (the fix) from datetime import datetime mock_org = MagicMock(spec=LiteLLM_OrganizationTable) mock_org.organization_id = "test-org-guardrails" mock_org.models = ["gpt-4", "gpt-3.5-turbo"] mock_org.budget_id = "budget-123" mock_org.created_by = "admin" mock_org.updated_by = "admin" mock_org.created_at = datetime(2024, 1, 1) mock_org.updated_at = datetime(2024, 1, 1) mock_org.litellm_budget_table = None mock_org.members = [] mock_org.teams = [] # Must be a list, not None mock_org.model_dump.return_value = { "organization_id": "test-org-guardrails", "models": ["gpt-4", "gpt-3.5-turbo"], "budget_id": "budget-123", "created_by": "admin", "updated_by": "admin", "created_at": datetime(2024, 1, 1), "updated_at": datetime(2024, 1, 1), "litellm_budget_table": None, "members": [], "teams": [], } with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ), patch( "litellm.proxy.proxy_server.premium_user", True # Required for guardrails feature ): # Mock existing team - must have compatible models with organization mock_existing_team = MagicMock() mock_existing_team.team_id = "team-guardrails-123" mock_existing_team.organization_id = None mock_existing_team.metadata = {} mock_existing_team.model_id = None mock_existing_team.models = ["gpt-4"] # Subset of org models to pass validation mock_existing_team.max_budget = None mock_existing_team.tpm_limit = None mock_existing_team.rpm_limit = None mock_existing_team.model_dump.return_value = { "team_id": "team-guardrails-123", "organization_id": None, "metadata": {}, "models": ["gpt-4"], "max_budget": None, "tpm_limit": None, "rpm_limit": None, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( return_value=mock_existing_team ) mock_cache.async_set_cache = AsyncMock() # Mock organization fetch - this is where the bug occurred # The fix ensures 'teams: True' is in the include clause mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock( return_value=mock_org ) # Mock team update mock_updated_team = MagicMock(spec=LiteLLM_TeamTable) mock_updated_team.team_id = "team-guardrails-123" mock_updated_team.organization_id = "test-org-guardrails" mock_updated_team.metadata = {"guardrails": ["aporia-pre-call", "aporia-post-call"]} mock_updated_team.litellm_model_table = None mock_updated_team.model_dump.return_value = { "team_id": "team-guardrails-123", "organization_id": "test-org-guardrails", "metadata": {"guardrails": ["aporia-pre-call", "aporia-post-call"]}, } mock_prisma.db.litellm_teamtable.update = AsyncMock( return_value=mock_updated_team ) mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) # Mock llm_router mock_router = MagicMock() with patch("litellm.proxy.proxy_server.llm_router", mock_router): # This should succeed without Pydantic validation error result = await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=org_admin_user, ) # Verify the team was updated successfully with guardrails assert result is not None assert result["data"].organization_id == "test-org-guardrails" assert result["data"].metadata["guardrails"] == ["aporia-pre-call", "aporia-post-call"] # Verify that organization fetch was called with proper include clause # The function is called twice: once by fetch_and_validate_organization (with include) # and once by get_org_object (without include). We verify the first call has 'teams'. assert mock_prisma.db.litellm_organizationtable.find_unique.call_count >= 1 # Get the first call (from fetch_and_validate_organization) first_call_kwargs = mock_prisma.db.litellm_organizationtable.find_unique.call_args_list[0].kwargs # Verify that 'teams' is included in the fetch assert "include" in first_call_kwargs assert "teams" in first_call_kwargs["include"] assert first_call_kwargs["include"]["teams"] is True def test_transform_teams_to_deleted_records(): from datetime import datetime, timezone user_api_key_dict = UserAPIKeyAuth( user_id="user-123", api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN.value, ) team1 = LiteLLM_TeamTable( team_id="team-1", team_alias="test-team-1", members_with_roles=[ Member(user_id="user-1", role="admin"), Member(user_id="user-2", role="user"), ], metadata={"test": "value"}, model_max_budget={}, model_spend={}, ) team2 = LiteLLM_TeamTable( team_id="team-2", team_alias="test-team-2", members_with_roles=[], metadata=None, model_max_budget={"gpt-4": {"budget_limit": 100.0}}, model_spend={}, ) records = _transform_teams_to_deleted_records( teams=[team1, team2], user_api_key_dict=user_api_key_dict, litellm_changed_by="admin-user", ) assert len(records) == 2 assert all("deleted_at" in record for record in records) assert all("deleted_by" in record for record in records) assert all("deleted_by_api_key" in record for record in records) assert all("litellm_changed_by" in record for record in records) assert all(record["deleted_by"] == "user-123" for record in records) # UserAPIKeyAuth hashes the api_key, so we check against the hashed value assert all(record["deleted_by_api_key"] == user_api_key_dict.api_key for record in records) assert all(record["litellm_changed_by"] == "admin-user" for record in records) record1 = records[0] assert record1["team_id"] == "team-1" assert isinstance(record1["members_with_roles"], str) assert isinstance(record1["metadata"], str) assert "litellm_model_table" not in record1 assert "object_permission" not in record1 assert "id" not in record1 record2 = records[1] assert record2["team_id"] == "team-2" # model_max_budget should be converted to JSON string if it exists if "model_max_budget" in record2: assert isinstance(record2["model_max_budget"], str) def test_transform_teams_to_deleted_records_empty_list(): user_api_key_dict = UserAPIKeyAuth( user_id="user-123", api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN.value, ) records = _transform_teams_to_deleted_records( teams=[], user_api_key_dict=user_api_key_dict, ) assert records == [] @pytest.mark.asyncio async def test_save_deleted_team_records(): mock_prisma_client = AsyncMock() mock_create_many = AsyncMock() mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many records = [ { "team_id": "team-1", "team_alias": "test-team-1", "deleted_at": "2024-01-01T00:00:00Z", "deleted_by": "admin", }, { "team_id": "team-2", "team_alias": "test-team-2", "deleted_at": "2024-01-01T00:00:00Z", "deleted_by": "admin", }, ] await _save_deleted_team_records(records=records, prisma_client=mock_prisma_client) mock_create_many.assert_called_once_with(data=records) @pytest.mark.asyncio async def test_save_deleted_team_records_empty_list(): mock_prisma_client = AsyncMock() mock_create_many = AsyncMock() mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many await _save_deleted_team_records(records=[], prisma_client=mock_prisma_client) mock_create_many.assert_not_called() @pytest.mark.asyncio async def test_persist_deleted_team_records(): mock_prisma_client = AsyncMock() mock_create_many = AsyncMock() mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many user_api_key_dict = UserAPIKeyAuth( user_id="user-123", api_key="sk-test", user_role=LitellmUserRoles.PROXY_ADMIN.value, ) team = LiteLLM_TeamTable( team_id="team-1", team_alias="test-team", members_with_roles=[ Member(user_id="user-1", role="admin"), ], metadata={}, model_max_budget={}, model_spend={}, ) await _persist_deleted_team_records( teams=[team], prisma_client=mock_prisma_client, user_api_key_dict=user_api_key_dict, litellm_changed_by="admin-user", ) mock_create_many.assert_called_once() call_args = mock_create_many.call_args assert "data" in call_args.kwargs records = call_args.kwargs["data"] assert len(records) == 1 assert records[0]["team_id"] == "team-1" assert records[0]["deleted_by"] == "user-123" assert records[0]["litellm_changed_by"] == "admin-user" @pytest.mark.asyncio async def test_delete_team_persists_deleted_teams(monkeypatch): from litellm.proxy._types import DeleteTeamRequest mock_prisma_client = AsyncMock() mock_user_api_key_dict = UserAPIKeyAuth( user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value, ) team1 = LiteLLM_TeamTable( team_id="team-1", team_alias="test-team-1", members_with_roles=[ Member(user_id="user-1", role="admin"), ], metadata={}, model_max_budget={}, model_spend={}, ) mock_find_unique = AsyncMock(return_value=team1) mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique mock_delete_data = AsyncMock(return_value={"deleted_teams": ["team-1"]}) mock_prisma_client.delete_data = mock_delete_data mock_create_many_teams = AsyncMock() mock_prisma_client.db.litellm_deletedteamtable.create_many = mock_create_many_teams mock_create_many_keys = AsyncMock() mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( mock_create_many_keys ) mock_find_many_keys = AsyncMock(return_value=[]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client, ) monkeypatch.setattr( "litellm.proxy.proxy_server.create_audit_log_for_update", AsyncMock(), ) monkeypatch.setattr( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin", ) monkeypatch.setattr( "litellm.proxy.management_endpoints.team_endpoints.team_member_delete", AsyncMock(return_value=team1), ) data = DeleteTeamRequest(team_ids=["team-1"]) result = await delete_team( data=data, http_request=MagicMock(), user_api_key_dict=mock_user_api_key_dict, litellm_changed_by="admin-user", ) mock_create_many_teams.assert_called_once() call_args = mock_create_many_teams.call_args assert "data" in call_args.kwargs records = call_args.kwargs["data"] assert len(records) == 1 assert records[0]["team_id"] == "team-1" assert records[0]["deleted_by"] == "admin-user" assert records[0]["litellm_changed_by"] == "admin-user" @pytest.mark.asyncio async def test_team_member_delete_persists_deleted_keys(monkeypatch): from litellm.proxy._types import TeamMemberDeleteRequest from litellm.proxy.management_endpoints.key_management_endpoints import ( LiteLLM_VerificationToken, ) mock_prisma_client = AsyncMock() mock_user_api_key_dict = UserAPIKeyAuth( user_id="admin-user", api_key="sk-admin", user_role=LitellmUserRoles.PROXY_ADMIN.value, ) team = LiteLLM_TeamTable( team_id="team-1", team_alias="test-team", members_with_roles=[ Member(user_id="user-123", role="admin"), ], metadata={}, model_max_budget={}, model_spend={}, ) key1 = LiteLLM_VerificationToken( token="hashed-token-1", user_id="user-123", team_id="team-1", key_alias="test-key-1", spend=100.0, max_budget=1000.0, models=["gpt-4"], aliases={}, config={}, permissions={}, metadata={}, model_max_budget={}, ) key2 = LiteLLM_VerificationToken( token="hashed-token-2", user_id="user-123", team_id="team-1", key_alias="test-key-2", spend=50.0, max_budget=500.0, models=["gpt-3.5-turbo"], aliases={}, config={}, permissions={}, metadata={}, model_max_budget={}, ) mock_find_unique_team = AsyncMock(return_value=team) mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique_team mock_find_many_user = AsyncMock( return_value=[ MagicMock( user_id="user-123", teams=["team-1"], model_dump=lambda: {"user_id": "user-123", "teams": ["team-1"]}, ) ] ) mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_user mock_update_team = AsyncMock() mock_prisma_client.db.litellm_teamtable.update = mock_update_team mock_update_user = AsyncMock() mock_prisma_client.db.litellm_usertable.update = mock_update_user mock_delete_membership = AsyncMock() mock_prisma_client.db.litellm_teammembership.delete_many = mock_delete_membership mock_find_many_keys = AsyncMock(return_value=[key1, key2]) mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys mock_delete_keys = AsyncMock() mock_prisma_client.db.litellm_verificationtoken.delete_many = mock_delete_keys mock_create_many_keys = AsyncMock() mock_prisma_client.db.litellm_deletedverificationtoken.create_many = ( mock_create_many_keys ) monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client, ) monkeypatch.setattr( "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", lambda **kwargs: True, ) data = TeamMemberDeleteRequest(team_id="team-1", user_id="user-123") result = await team_member_delete( data=data, user_api_key_dict=mock_user_api_key_dict, ) mock_create_many_keys.assert_called_once() call_args = mock_create_many_keys.call_args assert "data" in call_args.kwargs records = call_args.kwargs["data"] assert len(records) == 2 assert all(record["deleted_by"] == "admin-user" for record in records) assert all(record["team_id"] == "team-1" for record in records) assert all(record["user_id"] == "user-123" for record in records) mock_delete_keys.assert_called_once() @pytest.mark.asyncio async def test_new_team_negative_max_budget(): """ Test that NewTeamRequest model allows negative max_budget values. Validation is done at API level, not model level. This prevents GET requests from breaking when they receive data with negative budgets. """ from litellm.proxy._types import NewTeamRequest # Should not raise any errors at model level request = NewTeamRequest(team_alias="test-team", max_budget=-7.0) assert request.max_budget == -7.0 @pytest.mark.asyncio async def test_new_team_negative_team_member_budget(): """ Test that NewTeamRequest model allows negative team_member_budget values. Validation is done at API level, not model level. """ from litellm.proxy._types import NewTeamRequest # Should not raise any errors at model level request = NewTeamRequest(team_alias="test-team", team_member_budget=-10.0) assert request.team_member_budget == -10.0 @pytest.mark.asyncio async def test_update_team_negative_max_budget(): """ Test that UpdateTeamRequest model allows negative max_budget values. Validation is done at API level, not model level. """ from litellm.proxy._types import UpdateTeamRequest # Should not raise any errors at model level request = UpdateTeamRequest(team_id="test-team-id", max_budget=-5.0) assert request.max_budget == -5.0 @pytest.mark.asyncio async def test_update_team_negative_team_member_budget(): """ Test that UpdateTeamRequest model allows negative team_member_budget values. Validation is done at API level, not model level. """ from litellm.proxy._types import UpdateTeamRequest # Should not raise any errors at model level request = UpdateTeamRequest(team_id="test-team-id", team_member_budget=-15.0) assert request.team_member_budget == -15.0 # Parametrized tests for soft_budget in create endpoint @pytest.mark.parametrize( "soft_budget,max_budget,should_succeed,expected_soft_budget,expected_max_budget,error_message", [ # Test 1: Soft budget only - success + soft budget set (50.0, None, True, 50.0, None, None), # Test 2: Soft budget with higher max budget, success with both set (50.0, 100.0, True, 50.0, 100.0, None), # Test 3: Soft budget with lower max budget, fail (100.0, 50.0, False, None, None, "soft_budget (100.0) must be strictly lower than max_budget (50.0)"), # Test 4: Soft budget equal to max budget, fail (100.0, 100.0, False, None, None, "soft_budget (100.0) must be strictly lower than max_budget (100.0)"), ], ) @pytest.mark.asyncio async def test_new_team_soft_budget_validation( soft_budget, max_budget, should_succeed, expected_soft_budget, expected_max_budget, error_message ): """ Test soft_budget validation in /team/new endpoint. Covers: - Soft budget only - success + soft budget set - Soft budget with higher max budget, success with both set - Soft budget with lower max budget, fail """ from fastapi import Request from litellm.proxy._types import NewTeamRequest, ProxyException, UserAPIKeyAuth from litellm.proxy.management_endpoints.team_endpoints import new_team # Create admin user to bypass user budget checks admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user", models=[], ) # Create team request with soft_budget and optionally max_budget team_request = NewTeamRequest( team_alias="test-soft-budget-team", soft_budget=soft_budget, max_budget=max_budget, ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server._license_check" ) as mock_license, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), 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 from litellm.proxy._types import LiteLLM_UserTable mock_user_obj = LiteLLM_UserTable( user_id="admin-user", max_budget=None, # Admin has no budget limit ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Mock team creation mock_created_team = MagicMock() mock_created_team.team_id = "test-team-123" mock_created_team.team_alias = "test-soft-budget-team" mock_created_team.soft_budget = expected_soft_budget mock_created_team.max_budget = expected_max_budget mock_created_team.members_with_roles = [] mock_created_team.metadata = None mock_created_team.model_dump.return_value = { "team_id": "test-team-123", "team_alias": "test-soft-budget-team", "soft_budget": expected_soft_budget, "max_budget": expected_max_budget, "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 model table mock_prisma.db.litellm_modeltable = MagicMock() mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) # Mock user table operations mock_user = MagicMock() mock_user.user_id = "admin-user" mock_user.model_dump.return_value = {"user_id": "admin-user", "teams": ["test-team-123"]} 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 = { "team_id": "test-team-123", "user_id": "admin-user", "budget_id": None, } mock_prisma.db.litellm_teammembership = MagicMock() mock_prisma.db.litellm_teammembership.create = AsyncMock(return_value=mock_membership) if should_succeed: # Should NOT raise an exception result = await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=admin_user, ) # Verify the team was created successfully with correct values assert result is not None assert result["team_id"] == "test-team-123" if expected_soft_budget is not None: assert result["soft_budget"] == expected_soft_budget if expected_max_budget is not None: assert result["max_budget"] == expected_max_budget else: # Should raise ProxyException with pytest.raises(ProxyException) as exc_info: await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=admin_user, ) # Verify exception details assert exc_info.value.code == '400' if error_message: assert error_message in str(exc_info.value.message) # Parametrized tests for soft_budget in update endpoint @pytest.mark.parametrize( "existing_soft_budget,existing_max_budget,update_soft_budget,update_max_budget,should_succeed,expected_soft_budget,expected_max_budget,error_message", [ # Test 1: Soft budget only (no previous max_budget) - success with soft budget set (None, None, 50.0, None, True, 50.0, None, None), # Test 2: Soft budget with max budget - success if soft budget is strictly lower than max budget (None, None, 50.0, 100.0, True, 50.0, 100.0, None), # Test 3: Soft budget with max budget - fail if soft budget >= max budget (None, None, 100.0, 50.0, False, None, None, "soft_budget (100.0) must be strictly lower than max_budget (50.0)"), # Test 4: Only max budget with existing soft_budget, success with max_budget strictly greater (50.0, None, None, 100.0, True, 50.0, 100.0, None), # Test 5: Only max budget with existing soft_budget, fail if max_budget <= soft_budget (50.0, None, None, 50.0, False, None, None, "max_budget (50.0) must be strictly greater than soft_budget (50.0)"), # Test 6: Update both soft_budget and max_budget - success if soft < max (30.0, 100.0, 40.0, 80.0, True, 40.0, 80.0, None), # Test 7: Update both soft_budget and max_budget - fail if soft >= max (30.0, 100.0, 80.0, 40.0, False, None, None, "soft_budget (80.0) must be strictly lower than max_budget (40.0)"), ], ) @pytest.mark.asyncio async def test_update_team_soft_budget_validation( existing_soft_budget, existing_max_budget, update_soft_budget, update_max_budget, should_succeed, expected_soft_budget, expected_max_budget, error_message ): """ Test soft_budget validation in /team/update endpoint. Covers: - Soft budget only (no previous max_budget) - success with soft budget set - Soft budget with max budget - success if soft budget is strictly lower than max budget, fail otherwise - Only max budget with existing soft_budget, success with max_budget strictly greater, fail otherwise """ from fastapi import Request from litellm.proxy._types import ( LiteLLM_UserTable, ProxyException, UpdateTeamRequest, UserAPIKeyAuth, ) from litellm.proxy.management_endpoints.team_endpoints import update_team # Create admin user to bypass user budget checks admin_user = UserAPIKeyAuth( user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-user", models=[], ) # Create update request update_request = UpdateTeamRequest( team_id="test-team-123", soft_budget=update_soft_budget, max_budget=update_max_budget, ) dummy_request = MagicMock(spec=Request) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( "litellm.proxy.proxy_server.user_api_key_cache" ) as mock_cache, patch( "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" ), patch( "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() ) as mock_audit: # Mock existing team with existing budgets mock_existing_team = MagicMock() mock_existing_team.team_id = "test-team-123" mock_existing_team.organization_id = None mock_existing_team.soft_budget = existing_soft_budget mock_existing_team.max_budget = existing_max_budget mock_existing_team.model_dump.return_value = { "team_id": "test-team-123", "organization_id": None, "soft_budget": existing_soft_budget, "max_budget": existing_max_budget, } mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=mock_existing_team) # Mock user cache mock_user_obj = LiteLLM_UserTable( user_id="admin-user", max_budget=None, # Admin has no budget limit ) mock_cache.async_get_cache = AsyncMock(return_value=mock_user_obj) # Mock updated team - preserve existing values if not being updated final_soft_budget = update_soft_budget if update_soft_budget is not None else existing_soft_budget final_max_budget = update_max_budget if update_max_budget is not None else existing_max_budget mock_updated_team = MagicMock() mock_updated_team.team_id = "test-team-123" mock_updated_team.organization_id = None mock_updated_team.soft_budget = final_soft_budget mock_updated_team.max_budget = final_max_budget mock_updated_team.model_dump.return_value = { "team_id": "test-team-123", "organization_id": None, "soft_budget": final_soft_budget, "max_budget": final_max_budget, } mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=mock_updated_team) mock_prisma.jsonify_team_object = lambda db_data: db_data mock_cache.async_set_cache = AsyncMock() # Mock cache set for _cache_team_object if should_succeed: # Should NOT raise an exception result = await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=admin_user, ) # Verify the team was updated successfully with correct values assert result is not None assert result["data"].team_id == "test-team-123" # Verify soft_budget matches expected value (or final computed value if expected is None) if expected_soft_budget is not None: assert result["data"].soft_budget == expected_soft_budget else: assert result["data"].soft_budget == final_soft_budget # Verify max_budget matches expected value (or final computed value if expected is None) if expected_max_budget is not None: assert result["data"].max_budget == expected_max_budget else: assert result["data"].max_budget == final_max_budget else: # Should raise ProxyException with pytest.raises(ProxyException) as exc_info: await update_team( data=update_request, http_request=dummy_request, user_api_key_dict=admin_user, ) # Verify exception details assert exc_info.value.code == '400' if error_message: assert error_message in str(exc_info.value.message) @pytest.mark.asyncio async def test_new_team_positive_budgets_accepted(): """ Test that NewTeamRequest accepts positive budget values. """ from litellm.proxy._types import NewTeamRequest # Should not raise any errors request = NewTeamRequest( team_alias="test-team", max_budget=100.0, team_member_budget=50.0 ) assert request.max_budget == 100.0 assert request.team_member_budget == 50.0 @pytest.mark.asyncio async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth): """ Test that /team/new correctly handles router_settings by: 1. Accepting router_settings as a dict parameter 2. Serializing router_settings to JSON when saving to database 3. Storing router_settings in the team record """ # Configure mocked prisma client mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.get_data = AsyncMock(return_value=None) mock_db_client.update_data = AsyncMock(return_value=MagicMock()) mock_db_client.db = MagicMock() # Mock model table creation mock_db_client.db.litellm_modeltable = MagicMock() mock_db_client.db.litellm_modeltable.create = AsyncMock( return_value=MagicMock(id="model123") ) # Capture team table creation team_create_result = MagicMock( team_id="team-router-456", ) team_create_result.model_dump.return_value = { "team_id": "team-router-456", } mock_team_create = AsyncMock(return_value=team_create_result) mock_team_count = AsyncMock(return_value=0) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.create = mock_team_create mock_db_client.db.litellm_teamtable.count = mock_team_count mock_db_client.db.litellm_teamtable.update = AsyncMock( return_value=team_create_result ) # Mock user table mock_db_client.db.litellm_usertable = MagicMock() mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) from fastapi import Request from litellm.proxy._types import NewTeamRequest from litellm.proxy.management_endpoints.team_endpoints import new_team # Test router_settings with sample data router_settings_data = { "routing_strategy": "usage-based", "num_retries": 3, "retry_policy": {"max_retries": 5}, } # Build request with router_settings team_request = NewTeamRequest( team_alias="my-team-router", router_settings=router_settings_data, ) dummy_request = MagicMock(spec=Request) # Execute the endpoint function await new_team( data=team_request, http_request=dummy_request, user_api_key_dict=mock_admin_auth, ) # Verify team creation was called assert mock_team_create.call_count == 1 created_team_kwargs = mock_team_create.call_args.kwargs team_data = created_team_kwargs["data"] # Verify router_settings is serialized to JSON string assert "router_settings" in team_data assert isinstance(team_data["router_settings"], str) # Verify router_settings can be deserialized and matches input deserialized_settings = json.loads(team_data["router_settings"]) assert deserialized_settings == router_settings_data @pytest.mark.asyncio async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( mock_db_client, ): """ Test that non-team-admin users only see their own spend (filtered by their API keys) when calling /team/daily/activity endpoint. """ from litellm.proxy.management_endpoints.team_endpoints import ( get_team_daily_activity, ) # Create a non-admin user user_id = "test_user_123" team_id = "test_team_456" user_api_key_dict = UserAPIKeyAuth( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER ) # Mock user info mock_user_info = LiteLLM_UserTable( user_id=user_id, teams=[team_id], max_budget=1000.0, spend=0.0, user_email="test@example.com", user_role="internal_user", ) # Mock team with user as non-admin member mock_team_member = Member(user_id=user_id, role="user") mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.team_id = team_id mock_team.team_alias = "Test Team" mock_team.members_with_roles = [mock_team_member] mock_team.model_dump.return_value = { "team_id": team_id, "team_alias": "Test Team", "members_with_roles": [{"user_id": user_id, "role": "user"}], } # Mock user's API keys user_api_key_1 = MagicMock() user_api_key_1.token = "user_key_1" user_api_key_2 = MagicMock() user_api_key_2.token = "user_key_2" # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock( return_value=[mock_team] ) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( return_value=[user_api_key_1, user_api_key_2] ) # Mock get_user_object with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", new_callable=AsyncMock, ) as mock_get_user_object: mock_get_user_object.return_value = mock_user_info # Mock get_daily_activity to capture the api_key parameter with patch( "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", new_callable=AsyncMock, ) as mock_get_daily_activity: mock_get_daily_activity.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( team_ids=team_id, start_date="2024-01-01", end_date="2024-01-02", model=None, api_key=None, page=1, page_size=10, exclude_team_ids=None, user_api_key_dict=user_api_key_dict, ) # Verify get_daily_activity was called with user's API keys as filter mock_get_daily_activity.assert_called_once() call_kwargs = mock_get_daily_activity.call_args[1] assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"] assert call_kwargs["entity_id"] == [team_id] # Verify user's API keys were fetched mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() api_key_call_kwargs = ( mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] ) assert api_key_call_kwargs["where"] == {"user_id": user_id} @pytest.mark.asyncio async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client): """ Test that team admin users see all team spend (no API key filtering) when calling /team/daily/activity endpoint. """ from litellm.proxy.management_endpoints.team_endpoints import ( get_team_daily_activity, ) # Create a team admin user user_id = "test_admin_123" team_id = "test_team_456" user_api_key_dict = UserAPIKeyAuth( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER ) # Mock user info mock_user_info = LiteLLM_UserTable( user_id=user_id, teams=[team_id], max_budget=1000.0, spend=0.0, user_email="admin@example.com", user_role="internal_user", ) # Mock team with user as admin member mock_team_member = Member(user_id=user_id, role="admin") mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.team_id = team_id mock_team.team_alias = "Test Team" mock_team.members_with_roles = [mock_team_member] mock_team.model_dump.return_value = { "team_id": team_id, "team_alias": "Test Team", "members_with_roles": [{"user_id": user_id, "role": "admin"}], } # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock( return_value=[mock_team] ) # Mock get_user_object with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", new_callable=AsyncMock, ) as mock_get_user_object: mock_get_user_object.return_value = mock_user_info # Mock get_daily_activity to capture the api_key parameter with patch( "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", new_callable=AsyncMock, ) as mock_get_daily_activity: mock_get_daily_activity.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( team_ids=team_id, start_date="2024-01-01", end_date="2024-01-02", model=None, api_key=None, page=1, page_size=10, exclude_team_ids=None, user_api_key_dict=user_api_key_dict, ) # Verify get_daily_activity was called WITHOUT API key filtering mock_get_daily_activity.assert_called_once() call_kwargs = mock_get_daily_activity.call_args[1] assert call_kwargs["api_key"] is None assert call_kwargs["entity_id"] == [team_id] # Verify user's API keys were NOT fetched (since they're admin) if hasattr( mock_db_client.db.litellm_verificationtoken, "find_many" ) and mock_db_client.db.litellm_verificationtoken.find_many.called: # If it was called, that's unexpected for admin users assert False, "API keys should not be fetched for team admin users" @pytest.mark.asyncio async def test_get_team_daily_activity_member_with_permission_sees_all_spend( mock_db_client, ): """ Test that non-admin team members with /team/daily/activity permission can see all team spend (no API key filtering), same as team admins. """ from litellm.proxy.management_endpoints.team_endpoints import ( get_team_daily_activity, ) # Create a non-admin user user_id = "test_user_with_perm_123" team_id = "test_team_789" user_api_key_dict = UserAPIKeyAuth( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER ) # Mock user info mock_user_info = LiteLLM_UserTable( user_id=user_id, teams=[team_id], max_budget=1000.0, spend=0.0, user_email="member@example.com", user_role="internal_user", ) # Mock team with user as non-admin member AND /team/daily/activity permission mock_team_member = Member(user_id=user_id, role="user") mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.team_id = team_id mock_team.team_alias = "Test Team" mock_team.members_with_roles = [mock_team_member] mock_team.team_member_permissions = ["/team/daily/activity"] mock_team.model_dump.return_value = { "team_id": team_id, "team_alias": "Test Team", "members_with_roles": [{"user_id": user_id, "role": "user"}], "team_member_permissions": ["/team/daily/activity"], } # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock( return_value=[mock_team] ) # Mock get_user_object with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", new_callable=AsyncMock, ) as mock_get_user_object: mock_get_user_object.return_value = mock_user_info # Mock get_daily_activity to capture the api_key parameter with patch( "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", new_callable=AsyncMock, ) as mock_get_daily_activity: mock_get_daily_activity.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( team_ids=team_id, start_date="2024-01-01", end_date="2024-01-02", model=None, api_key=None, page=1, page_size=10, exclude_team_ids=None, user_api_key_dict=user_api_key_dict, ) # Verify get_daily_activity was called WITHOUT API key filtering mock_get_daily_activity.assert_called_once() call_kwargs = mock_get_daily_activity.call_args[1] assert call_kwargs["api_key"] is None assert call_kwargs["entity_id"] == [team_id] # Verify user's API keys were NOT fetched if hasattr( mock_db_client.db.litellm_verificationtoken, "find_many" ) and mock_db_client.db.litellm_verificationtoken.find_many.called: assert ( False ), "API keys should not be fetched for members with /team/daily/activity permission" @pytest.mark.asyncio async def test_get_team_daily_activity_member_without_permission_filters_by_keys( mock_db_client, ): """ Test that non-admin team members WITHOUT /team/daily/activity permission still have their results filtered by their own API keys. """ from litellm.proxy.management_endpoints.team_endpoints import ( get_team_daily_activity, ) # Create a non-admin user user_id = "test_user_no_perm_123" team_id = "test_team_789" user_api_key_dict = UserAPIKeyAuth( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER ) # Mock user info mock_user_info = LiteLLM_UserTable( user_id=user_id, teams=[team_id], max_budget=1000.0, spend=0.0, user_email="member@example.com", user_role="internal_user", ) # Mock team with user as non-admin member and NO usage permission mock_team_member = Member(user_id=user_id, role="user") mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.team_id = team_id mock_team.team_alias = "Test Team" mock_team.members_with_roles = [mock_team_member] mock_team.team_member_permissions = ["/key/info"] mock_team.model_dump.return_value = { "team_id": team_id, "team_alias": "Test Team", "members_with_roles": [{"user_id": user_id, "role": "user"}], "team_member_permissions": ["/key/info"], } # Mock user's API keys user_api_key_1 = MagicMock() user_api_key_1.token = "user_key_abc" user_api_key_2 = MagicMock() user_api_key_2.token = "user_key_def" # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock( return_value=[mock_team] ) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( return_value=[user_api_key_1, user_api_key_2] ) # Mock get_user_object with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", new_callable=AsyncMock, ) as mock_get_user_object: mock_get_user_object.return_value = mock_user_info # Mock get_daily_activity to capture the api_key parameter with patch( "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", new_callable=AsyncMock, ) as mock_get_daily_activity: mock_get_daily_activity.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( team_ids=team_id, start_date="2024-01-01", end_date="2024-01-02", model=None, api_key=None, page=1, page_size=10, exclude_team_ids=None, user_api_key_dict=user_api_key_dict, ) # Verify get_daily_activity was called WITH API key filtering mock_get_daily_activity.assert_called_once() call_kwargs = mock_get_daily_activity.call_args[1] assert call_kwargs["api_key"] == ["user_key_abc", "user_key_def"] assert call_kwargs["entity_id"] == [team_id] # Verify user's API keys were fetched mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() @pytest.mark.asyncio async def test_update_team_with_router_settings(mock_db_client, mock_admin_auth): """ Test that /team/update correctly handles router_settings by: 1. Accepting router_settings as a dict parameter 2. Serializing router_settings to JSON when updating database 3. Updating router_settings in the team record """ # Configure mocked prisma client mock_db_client.jsonify_team_object = lambda db_data: db_data mock_db_client.db = MagicMock() # Mock existing team row existing_team_mock = MagicMock() existing_team_mock.team_id = "team-router-update-789" existing_team_mock.organization_id = None existing_team_mock.models = [] existing_team_mock.members_with_roles = [] existing_team_mock.model_dump.return_value = { "team_id": "team-router-update-789", "organization_id": None, "models": [], "members_with_roles": [], } # Mock team table find_unique and update updated_team_result = MagicMock( team_id="team-router-update-789", ) updated_team_result.model_dump.return_value = { "team_id": "team-router-update-789", } mock_team_find_unique = AsyncMock(return_value=existing_team_mock) mock_team_update = AsyncMock(return_value=updated_team_result) mock_db_client.db.litellm_teamtable = MagicMock() mock_db_client.db.litellm_teamtable.find_unique = mock_team_find_unique mock_db_client.db.litellm_teamtable.update = mock_team_update from fastapi import Request from litellm.proxy._types import UpdateTeamRequest from litellm.proxy.management_endpoints.team_endpoints import update_team # Test router_settings with updated data router_settings_data = { "routing_strategy": "latency-based", "num_retries": 2, } # Build update request with router_settings team_update_request = UpdateTeamRequest( team_id="team-router-update-789", router_settings=router_settings_data, ) dummy_request = MagicMock(spec=Request) # Execute the endpoint function await update_team( data=team_update_request, http_request=dummy_request, user_api_key_dict=mock_admin_auth, ) # Verify team update was called assert mock_team_update.call_count == 1 updated_team_kwargs = mock_team_update.call_args.kwargs team_data = updated_team_kwargs["data"] # Verify router_settings is serialized to JSON string assert "router_settings" in team_data assert isinstance(team_data["router_settings"], str) # Verify router_settings can be deserialized and matches input deserialized_settings = json.loads(team_data["router_settings"]) assert deserialized_settings == router_settings_data @pytest.mark.asyncio async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( mock_db_client, ): """ Test that non-team-admin users only see their own spend (filtered by their API keys) when calling /team/daily/activity endpoint. """ from litellm.proxy.management_endpoints.team_endpoints import ( get_team_daily_activity, ) # Create a non-admin user user_id = "test_user_123" team_id = "test_team_456" user_api_key_dict = UserAPIKeyAuth( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER ) # Mock user info mock_user_info = LiteLLM_UserTable( user_id=user_id, teams=[team_id], max_budget=1000.0, spend=0.0, user_email="test@example.com", user_role="internal_user", ) # Mock team with user as non-admin member mock_team_member = Member(user_id=user_id, role="user") mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.team_id = team_id mock_team.team_alias = "Test Team" mock_team.members_with_roles = [mock_team_member] mock_team.model_dump.return_value = { "team_id": team_id, "team_alias": "Test Team", "members_with_roles": [{"user_id": user_id, "role": "user"}], } # Mock user's API keys user_api_key_1 = MagicMock() user_api_key_1.token = "user_key_1" user_api_key_2 = MagicMock() user_api_key_2.token = "user_key_2" # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock( return_value=[mock_team] ) mock_db_client.db.litellm_verificationtoken.find_many = AsyncMock( return_value=[user_api_key_1, user_api_key_2] ) # Mock get_user_object with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", new_callable=AsyncMock, ) as mock_get_user_object: mock_get_user_object.return_value = mock_user_info # Mock get_daily_activity to capture the api_key parameter with patch( "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", new_callable=AsyncMock, ) as mock_get_daily_activity: mock_get_daily_activity.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( team_ids=team_id, start_date="2024-01-01", end_date="2024-01-02", model=None, api_key=None, page=1, page_size=10, exclude_team_ids=None, user_api_key_dict=user_api_key_dict, ) # Verify get_daily_activity was called with user's API keys as filter mock_get_daily_activity.assert_called_once() call_kwargs = mock_get_daily_activity.call_args[1] assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"] assert call_kwargs["entity_id"] == [team_id] # Verify user's API keys were fetched mock_db_client.db.litellm_verificationtoken.find_many.assert_called_once() api_key_call_kwargs = ( mock_db_client.db.litellm_verificationtoken.find_many.call_args[1] ) assert api_key_call_kwargs["where"] == {"user_id": user_id} @pytest.mark.asyncio async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client): """ Test that team admin users see all team spend (no API key filtering) when calling /team/daily/activity endpoint. """ from litellm.proxy.management_endpoints.team_endpoints import ( get_team_daily_activity, ) # Create a team admin user user_id = "test_admin_123" team_id = "test_team_456" user_api_key_dict = UserAPIKeyAuth( user_id=user_id, user_role=LitellmUserRoles.INTERNAL_USER ) # Mock user info mock_user_info = LiteLLM_UserTable( user_id=user_id, teams=[team_id], max_budget=1000.0, spend=0.0, user_email="admin@example.com", user_role="internal_user", ) # Mock team with user as admin member mock_team_member = Member(user_id=user_id, role="admin") mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.team_id = team_id mock_team.team_alias = "Test Team" mock_team.members_with_roles = [mock_team_member] mock_team.model_dump.return_value = { "team_id": team_id, "team_alias": "Test Team", "members_with_roles": [{"user_id": user_id, "role": "admin"}], } # Setup mocks mock_db_client.db.litellm_teamtable.find_many = AsyncMock( return_value=[mock_team] ) # Mock get_user_object with patch( "litellm.proxy.management_endpoints.team_endpoints.get_user_object", new_callable=AsyncMock, ) as mock_get_user_object: mock_get_user_object.return_value = mock_user_info # Mock get_daily_activity to capture the api_key parameter with patch( "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", new_callable=AsyncMock, ) as mock_get_daily_activity: mock_get_daily_activity.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( team_ids=team_id, start_date="2024-01-01", end_date="2024-01-02", model=None, api_key=None, page=1, page_size=10, exclude_team_ids=None, user_api_key_dict=user_api_key_dict, ) # Verify get_daily_activity was called WITHOUT API key filtering mock_get_daily_activity.assert_called_once() call_kwargs = mock_get_daily_activity.call_args[1] assert call_kwargs["api_key"] is None assert call_kwargs["entity_id"] == [team_id] # Verify user's API keys were NOT fetched (since they're admin) if hasattr( mock_db_client.db.litellm_verificationtoken, "find_many" ) and mock_db_client.db.litellm_verificationtoken.find_many.called: # If it was called, that's unexpected for admin users assert False, "API keys should not be fetched for team admin users" @pytest.mark.asyncio async def test_validate_and_populate_member_user_info_both_provided_match(): """ Test _validate_and_populate_member_user_info when both user_email and user_id are provided and they match the same user in the database. """ # Create member with both user_email and user_id member = Member(user_email="test@example.com", user_id="user-123", role="user") # Mock prisma client mock_prisma_client = MagicMock() # Mock user object that matches both email and user_id mock_user = MagicMock() mock_user.user_id = "user-123" mock_user.user_email = "test@example.com" # Mock get_data to return single user matching email mock_prisma_client.get_data = AsyncMock(return_value=[mock_user]) # Call the function result = await _validate_and_populate_member_user_info( member=member, prisma_client=mock_prisma_client, ) # Verify result matches input (both already provided and match) assert result.user_email == "test@example.com" assert result.user_id == "user-123" # Verify get_data was called with correct parameters mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": "test@example.com"}, table_name="user", query_type="find_all", ) @pytest.mark.asyncio async def test_validate_and_populate_member_user_info_only_email_provided(): """ Test _validate_and_populate_member_user_info when only user_email is provided. Should populate user_id from database. """ # Create member with only user_email member = Member(user_email="test@example.com", user_id=None, role="user") # Mock prisma client mock_prisma_client = MagicMock() # Mock user object from find_first mock_user_find_first = MagicMock() mock_user_find_first.user_id = "user-456" mock_user_find_first.user_email = "test@example.com" # Mock find_first to return the user mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( return_value=mock_user_find_first ) # Mock get_data to return single user (no duplicates) mock_prisma_client.get_data = AsyncMock(return_value=[mock_user_find_first]) # Call the function result = await _validate_and_populate_member_user_info( member=member, prisma_client=mock_prisma_client, ) # Verify user_id was populated assert result.user_email == "test@example.com" assert result.user_id == "user-456" # Verify find_first was called with correct parameters mock_prisma_client.db.litellm_usertable.find_first.assert_called_once_with( where={"user_email": {"equals": "test@example.com", "mode": "insensitive"}} ) # Verify get_data was called to check for duplicates mock_prisma_client.get_data.assert_called_once_with( key_val={"user_email": "test@example.com"}, table_name="user", query_type="find_all", ) @pytest.mark.asyncio async def test_validate_and_populate_member_user_info_only_user_id_not_found(): """ Test _validate_and_populate_member_user_info when only user_id is provided but the user doesn't exist in the database. Should allow it to pass with user_email as None (will be upserted later). """ # Create member with only user_id member = Member(user_email=None, user_id="nonexistent-user", role="user") # Mock prisma client mock_prisma_client = MagicMock() # Mock find_unique to return None (user not found) mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) # Call the function - should NOT raise an exception result = await _validate_and_populate_member_user_info( member=member, prisma_client=mock_prisma_client, ) # Verify the result - should return member with user_id set and user_email as None assert result.user_id == "nonexistent-user" assert result.user_email is None assert result.role == "user" # Verify find_unique was called with correct parameters mock_prisma_client.db.litellm_usertable.find_unique.assert_called_once_with( where={"user_id": "nonexistent-user"} ) @pytest.mark.asyncio async def test_list_available_teams_returns_empty_list_when_none_configured(): """ Test that /team/available returns an empty list when no available teams are configured, instead of raising an exception. """ import litellm mock_request = MagicMock() mock_user_key = UserAPIKeyAuth(user_id="test-user", token="fake-token") with patch( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client ): # Case 1: default_internal_user_params is None original = litellm.default_internal_user_params litellm.default_internal_user_params = None result = await list_available_teams( http_request=mock_request, user_api_key_dict=mock_user_key, ) assert result == [] # Case 2: default_internal_user_params exists but has no "available_teams" key litellm.default_internal_user_params = {"some_other_param": "value"} result = await list_available_teams( http_request=mock_request, user_api_key_dict=mock_user_key, ) assert result == [] litellm.default_internal_user_params = original