diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a8184a34d45..b84eaf45922 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -20,6 +20,7 @@ from litellm.proxy._types import ( LiteLLM_TeamTableCachedObj, LiteLLM_VerificationToken, LitellmUserRoles, + Member, ProxyException, UpdateKeyRequest, ) @@ -29,6 +30,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_team_key_limits, _common_key_generation_helper, _list_key_helper, + can_delete_verification_token, check_org_key_model_specific_limits, check_team_key_model_specific_limits, generate_key_helper_fn, @@ -2613,3 +2615,383 @@ def test_check_org_key_model_specific_limits_org_model_tpm_overallocation(): "Allocated TPM limit=17000 + Key TPM limit=4000 is greater than organization TPM limit=20000" in str(exc_info.value.detail) ) + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_proxy_admin_team_key(monkeypatch): + """Test that proxy admin can delete any team key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_proxy_admin_personal_key(monkeypatch): + """Test that proxy admin can delete any personal key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin-user", + api_key="sk-admin", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_team_admin_own_team(monkeypatch): + """Test that team admin can delete team keys from their own team.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="team-admin-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="team-admin-user", role="admin"), + Member(user_id="other-user", role="user"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_team_admin_different_team(monkeypatch): + """Test that team admin cannot delete team keys from a different team.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="other-user", + team_id="test-team-456", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="team-admin-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-456", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="different-admin", role="admin"), + Member(user_id="other-user", role="user"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_key_owner_team_key(monkeypatch): + """Test that key owner can delete their own team key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="key-owner-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="key-owner-user", role="user"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_key_owner_personal_key(monkeypatch): + """Test that key owner can delete their own personal key.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="key-owner-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is True + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_other_user_team_key(monkeypatch): + """Test that other user cannot delete team keys they don't own and aren't admin for.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id="test-team-123", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="other-user", + api_key="sk-user", + ) + + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=None, + rpm_limit=None, + max_budget=None, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[ + Member(user_id="key-owner-user", role="user"), + Member(user_id="other-user", role="user"), + Member(user_id="team-admin-user", role="admin"), + ], + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return team_table + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_other_user_personal_key(monkeypatch): + """Test that other user cannot delete personal keys they don't own.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="other-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_team_key_no_team_found(monkeypatch): + """Test that deletion fails when team is not found in database.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id="key-owner-user", + team_id="non-existent-team", + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="key-owner-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + async def mock_get_team_object(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_team_object", + mock_get_team_object, + ) + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False + + +@pytest.mark.asyncio +async def test_can_delete_verification_token_personal_key_no_user_id(monkeypatch): + """Test that deletion fails for personal key when key has no user_id.""" + key_info = LiteLLM_VerificationToken( + token="test-token", + user_id=None, + team_id=None, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + user_id="some-user", + api_key="sk-user", + ) + + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + + result = await can_delete_verification_token( + key_info=key_info, + user_api_key_cache=mock_user_api_key_cache, + user_api_key_dict=user_api_key_dict, + prisma_client=mock_prisma_client, + ) + + assert result is False