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 e3aa7d58872..35ecdd8e0ae 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 @@ -15,12 +15,14 @@ from fastapi import HTTPException from litellm.proxy._types import ( GenerateKeyRequest, + LiteLLM_TeamTableCachedObj, LiteLLM_VerificationToken, LitellmUserRoles, UpdateKeyRequest, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.proxy.management_endpoints.key_management_endpoints import ( + _check_team_key_limits, _common_key_generation_helper, _list_key_helper, generate_key_helper_fn, @@ -1040,7 +1042,7 @@ async def test_unblock_key_invalid_key_format(monkeypatch): def test_validate_key_team_change_with_member_permissions(): """ Test validate_key_team_change function with team member permissions. - + This test covers the new logic that allows team members with specific permissions to update keys, not just team admins. """ @@ -1054,111 +1056,107 @@ def test_validate_key_team_change_with_member_permissions(): mock_key.models = ["gpt-4"] mock_key.tpm_limit = None mock_key.rpm_limit = None - + mock_team = MagicMock() - mock_team.team_id = "test-team-456" + mock_team.team_id = "test-team-456" mock_team.members_with_roles = [] mock_team.tpm_limit = None mock_team.rpm_limit = None - + mock_change_initiator = MagicMock() mock_change_initiator.user_id = "test-user-123" - + mock_router = MagicMock() - + # Mock the member object returned by _get_user_in_team mock_member_object = MagicMock() - - with patch('litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model'): - with patch('litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team') as mock_get_user: - with patch('litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin') as mock_is_admin: - with patch('litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint') as mock_has_perms: - + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.can_team_access_model" + ): + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._get_user_in_team" + ) as mock_get_user: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._is_user_team_admin" + ) as mock_is_admin: + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.does_team_member_have_permissions_for_endpoint" + ) as mock_has_perms: + mock_get_user.return_value = mock_member_object mock_is_admin.return_value = False mock_has_perms.return_value = True - + # This should not raise an exception due to member permissions validate_key_team_change( key=mock_key, team=mock_team, change_initiated_by=mock_change_initiator, - llm_router=mock_router + llm_router=mock_router, ) - + # Verify the permission check was called with correct parameters mock_has_perms.assert_called_once_with( team_member_object=mock_member_object, team_table=mock_team, - route=KeyManagementRoutes.KEY_UPDATE.value + route=KeyManagementRoutes.KEY_UPDATE.value, ) def test_key_rotation_fields_helper(): """ Test the key data update logic for rotation fields. - + This test focuses on the core logic that adds rotation fields to key_data when auto_rotate is enabled, without the complexity of full key generation. """ # Test Case 1: With rotation enabled - key_data = { - "models": ["gpt-3.5-turbo"], - "user_id": "test-user" - } - + key_data = {"models": ["gpt-3.5-turbo"], "user_id": "test-user"} + auto_rotate = True rotation_interval = "30d" - + # Simulate the rotation logic from generate_key_helper_fn if auto_rotate and rotation_interval: - key_data.update({ - "auto_rotate": auto_rotate, - "rotation_interval": rotation_interval - }) - + key_data.update( + {"auto_rotate": auto_rotate, "rotation_interval": rotation_interval} + ) + # Verify rotation fields are added assert key_data["auto_rotate"] == True assert key_data["rotation_interval"] == "30d" assert key_data["models"] == ["gpt-3.5-turbo"] # Original fields preserved - + # Test Case 2: Without rotation enabled - key_data2 = { - "models": ["gpt-4"], - "user_id": "test-user" - } - + key_data2 = {"models": ["gpt-4"], "user_id": "test-user"} + auto_rotate2 = False rotation_interval2 = None - + # Simulate the rotation logic if auto_rotate2 and rotation_interval2: - key_data2.update({ - "auto_rotate": auto_rotate2, - "rotation_interval": rotation_interval2 - }) - + key_data2.update( + {"auto_rotate": auto_rotate2, "rotation_interval": rotation_interval2} + ) + # Verify rotation fields are NOT added assert "auto_rotate" not in key_data2 assert "rotation_interval" not in key_data2 assert key_data2["models"] == ["gpt-4"] # Original fields preserved - + # Test Case 3: auto_rotate=True but no interval - key_data3 = { - "models": ["claude-3"], - "user_id": "test-user" - } - + key_data3 = {"models": ["claude-3"], "user_id": "test-user"} + auto_rotate3 = True rotation_interval3 = None - + # Simulate the rotation logic if auto_rotate3 and rotation_interval3: - key_data3.update({ - "auto_rotate": auto_rotate3, - "rotation_interval": rotation_interval3 - }) - + key_data3.update( + {"auto_rotate": auto_rotate3, "rotation_interval": rotation_interval3} + ) + # Verify rotation fields are NOT added (missing interval) assert "auto_rotate" not in key_data3 assert "rotation_interval" not in key_data3 @@ -1181,27 +1179,24 @@ async def test_update_key_fn_auto_rotate_enable(): team_id=None, auto_rotate=False, rotation_interval=None, - metadata={} + metadata={}, ) - + # Test enabling auto rotation update_request = UpdateKeyRequest( - key="test-token", - auto_rotate=True, - rotation_interval="30d" + key="test-token", auto_rotate=True, rotation_interval="30d" ) - + result = await prepare_key_update_data( - data=update_request, - existing_key_row=existing_key + data=update_request, existing_key_row=existing_key ) - + # Verify rotation fields are included assert result["auto_rotate"] is True assert result["rotation_interval"] == "30d" -@pytest.mark.asyncio +@pytest.mark.asyncio async def test_update_key_fn_auto_rotate_disable(): """Test that update_key_fn properly handles disabling auto rotation.""" from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest @@ -1218,19 +1213,400 @@ async def test_update_key_fn_auto_rotate_disable(): team_id=None, auto_rotate=True, rotation_interval="30d", - metadata={} + metadata={}, ) - + # Test disabling auto rotation - update_request = UpdateKeyRequest( - key="test-token", - auto_rotate=False - ) - + update_request = UpdateKeyRequest(key="test-token", auto_rotate=False) + result = await prepare_key_update_data( - data=update_request, - existing_key_row=existing_key + data=update_request, existing_key_row=existing_key ) - + # Verify auto_rotate is set to False assert result["auto_rotate"] is False + + +@pytest.mark.asyncio +async def test_check_team_key_limits_no_existing_keys(): + """ + Test _check_team_key_limits when team has no existing keys. + Should allow any TPM/RPM limits within team bounds. + """ + # Mock prisma client + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-123", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with limits within team bounds + data = GenerateKeyRequest( + tpm_limit=5000, + rpm_limit=500, + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + # Verify database was queried + mock_prisma_client.db.litellm_verificationtoken.find_many.assert_called_once_with( + where={"team_id": "test-team-123"} + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_with_existing_keys_within_bounds(): + """ + Test _check_team_key_limits when team has existing keys but total allocation + is still within team limits. + """ + # Create mock existing keys + existing_key1 = MagicMock() + existing_key1.tpm_limit = 3000 + existing_key1.rpm_limit = 200 + + existing_key2 = MagicMock() + existing_key2.tpm_limit = 2000 + existing_key2.rpm_limit = 300 + + existing_key3 = MagicMock() + existing_key3.tpm_limit = None # Should be ignored in calculation + existing_key3.rpm_limit = None # Should be ignored in calculation + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2, existing_key3] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-456", + team_alias="test-team", + tpm_limit=10000, # Total: 3000 + 2000 + 4000 (new) = 9000 < 10000 ✓ + rpm_limit=1000, # Total: 200 + 300 + 400 (new) = 900 < 1000 ✓ + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that would still be within bounds + data = GenerateKeyRequest( + tpm_limit=4000, + rpm_limit=400, + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_tpm_overallocation(): + """ + Test _check_team_key_limits when new key would cause TPM overallocation. + Should raise HTTPException with appropriate error message. + """ + # Create mock existing keys with high TPM usage + existing_key1 = MagicMock() + existing_key1.tpm_limit = 6000 + existing_key1.rpm_limit = 100 + + existing_key2 = MagicMock() + existing_key2.tpm_limit = 3000 + existing_key2.rpm_limit = 200 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-789", + team_alias="test-team", + tpm_limit=10000, # Allocated: 6000 + 3000 = 9000, New: 2000, Total: 11000 > 10000 ✗ + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that would exceed TPM limits + data = GenerateKeyRequest( + tpm_limit=2000, + rpm_limit=100, + ) + + # Should raise HTTPException for TPM overallocation + with pytest.raises(HTTPException) as exc_info: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert ( + "Allocated TPM limit=9000 + Key TPM limit=2000 is greater than team TPM limit=10000" + in str(exc_info.value.detail) + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_rpm_overallocation(): + """ + Test _check_team_key_limits when new key would cause RPM overallocation. + Should raise HTTPException with appropriate error message. + """ + # Create mock existing keys with high RPM usage + existing_key1 = MagicMock() + existing_key1.tpm_limit = 1000 + existing_key1.rpm_limit = 600 + + existing_key2 = MagicMock() + existing_key2.tpm_limit = 2000 + existing_key2.rpm_limit = 300 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-101", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, # Allocated: 600 + 300 = 900, New: 200, Total: 1100 > 1000 ✗ + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that would exceed RPM limits + data = GenerateKeyRequest( + tpm_limit=1000, + rpm_limit=200, + ) + + # Should raise HTTPException for RPM overallocation + with pytest.raises(HTTPException) as exc_info: + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.status_code == 400 + assert ( + "Allocated RPM limit=900 + Key RPM limit=200 is greater than team RPM limit=1000" + in str(exc_info.value.detail) + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_no_team_limits(): + """ + Test _check_team_key_limits when team has no TPM/RPM limits set. + Should allow any key limits since there are no team constraints. + """ + # Create mock existing keys + existing_key = MagicMock() + existing_key.tpm_limit = 5000 + existing_key.rpm_limit = 500 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + # Create team table with no limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-202", + team_alias="test-team", + tpm_limit=None, # No team limit + rpm_limit=None, # No team limit + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with any limits + data = GenerateKeyRequest( + tpm_limit=10000, # High limit should be allowed + rpm_limit=2000, # High limit should be allowed + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_no_key_limits(): + """ + Test _check_team_key_limits when new key has no TPM/RPM limits. + Should not raise any exceptions since no limits are being allocated. + """ + # Create mock existing keys + existing_key = MagicMock() + existing_key.tpm_limit = 8000 + existing_key.rpm_limit = 800 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-303", + team_alias="test-team", + tpm_limit=10000, + rpm_limit=1000, + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with no limits + data = GenerateKeyRequest( + tpm_limit=None, # No limit being set + rpm_limit=None, # No limit being set + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_mixed_scenarios(): + """ + Test _check_team_key_limits with mixed scenarios: + - Some existing keys have limits, others don't + - New key has only one type of limit + - Team has only one type of limit + """ + # Create mock existing keys with mixed limits + existing_key1 = MagicMock() + existing_key1.tpm_limit = 3000 + existing_key1.rpm_limit = None # No RPM limit + + existing_key2 = MagicMock() + existing_key2.tpm_limit = None # No TPM limit + existing_key2.rpm_limit = 400 + + existing_key3 = MagicMock() + existing_key3.tpm_limit = 2000 + existing_key3.rpm_limit = 300 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key1, existing_key2, existing_key3] + ) + + # Create team table with only TPM limit + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-404", + team_alias="test-team", + tpm_limit=10000, # Allocated: 3000 + 0 + 2000 = 5000, New: 4000, Total: 9000 < 10000 ✓ + rpm_limit=None, # No team RPM limit + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request with only TPM limit + data = GenerateKeyRequest( + tpm_limit=4000, + rpm_limit=None, # No RPM limit being set + ) + + # Should not raise any exception + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + ) + + +@pytest.mark.asyncio +async def test_check_team_key_limits_exact_boundary(): + """ + Test _check_team_key_limits when allocation exactly matches team limits. + Should allow the allocation (boundary case). + """ + # Create mock existing keys + existing_key = MagicMock() + existing_key.tpm_limit = 7000 + existing_key.rpm_limit = 700 + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[existing_key] + ) + + # Create team table with limits + team_table = LiteLLM_TeamTableCachedObj( + team_id="test-team-505", + team_alias="test-team", + tpm_limit=10000, # Allocated: 7000, New: 3000, Total: 10000 = 10000 ✓ + rpm_limit=1000, # Allocated: 700, New: 300, Total: 1000 = 1000 ✓ + max_budget=100.0, + spend=0.0, + models=[], + blocked=False, + members_with_roles=[], + ) + + # Create request that exactly matches remaining capacity + data = GenerateKeyRequest( + tpm_limit=3000, + rpm_limit=300, + ) + + # Should not raise any exception (exact boundary should be allowed) + await _check_team_key_limits( + team_table=team_table, + data=data, + prisma_client=mock_prisma_client, + )