import pytest from fastapi import HTTPException from starlette.requests import Request import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import TeamMemberUpdateRequest from litellm.proxy.management_endpoints.team_endpoints import team_member_update @pytest.mark.asyncio async def test_team_member_update_admin_requires_premium(monkeypatch): # Arrange: patch prisma_client and premium_user monkeypatch.setattr(proxy_server, "prisma_client", object()) monkeypatch.setattr(proxy_server, "premium_user", False) # Create a request body that tries to set role=admin data = TeamMemberUpdateRequest( team_id="team-1234", user_id="user-1", user_email=None, role="admin", max_budget_in_team=None, ) scope = {"type": "http", "method": "POST", "path": "/team/member_update"} request = Request(scope) # We don't need a full auth object since premium check happens before auth is used auth = object() # Act & Assert: expect HTTPException 400 with the exact premium feature message with pytest.raises(HTTPException) as exc_info: await team_member_update(data, request, auth) assert exc_info.value.status_code == 400 expected_msg = ( "Assigning team admins is a premium feature. You must be a LiteLLM Enterprise user to use this feature. " "If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/#trial. " "Pricing: https://www.litellm.ai/#pricing" ) assert exc_info.value.detail == expected_msg