From 0c8b311155f01f5075ca4deb09ac3a5b390ad2c6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 11 Sep 2025 19:15:15 -0700 Subject: [PATCH] test: add unit testing for both flows on key unblock --- .../test_key_management_endpoints.py | 245 +++++++++++++++--- 1 file changed, 202 insertions(+), 43 deletions(-) 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 3a597adef06..2747fab77fa 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 @@ -183,7 +183,9 @@ async def test_budget_reset_and_expires_at_first_of_month(monkeypatch): assert ( response_date.month == expected_month ), f"Expected month {expected_month}, got {response_date.month} for {key}" - assert response_date.day == 1, f"Expected day 1, got {response_date.day} for {key}" + assert ( + response_date.day == 1 + ), f"Expected day 1, got {response_date.day} for {key}" @pytest.mark.asyncio @@ -507,7 +509,6 @@ def test_get_new_token_with_invalid_key(): assert "New key must start with 'sk-'" in str(exc_info.value.detail) - @pytest.mark.asyncio async def test_generate_service_account_requires_team_id(): with pytest.raises(HTTPException): @@ -529,11 +530,12 @@ async def test_generate_service_account_works_with_team_id(): from unittest.mock import patch # Mock the database and router dependencies from proxy_server - with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma, \ - patch('litellm.proxy.proxy_server.llm_router') as mock_router, \ - patch('litellm.proxy.proxy_server.premium_user', False), \ - patch('litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn') as mock_generate_key: - + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch( + "litellm.proxy.proxy_server.llm_router" + ) as mock_router, patch("litellm.proxy.proxy_server.premium_user", False), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key: + # Configure mocks mock_prisma.return_value = AsyncMock() mock_router.return_value = None @@ -542,9 +544,9 @@ async def test_generate_service_account_works_with_team_id(): "key": "sk-test-key", "expires": None, "user_id": "test-user", - "team_id": "IJ" + "team_id": "IJ", } - + # This should not raise an exception since team_id is provided await _common_key_generation_helper( data=GenerateKeyRequest( @@ -559,7 +561,6 @@ async def test_generate_service_account_works_with_team_id(): ) - @pytest.mark.asyncio async def test_update_service_account_requires_team_id(): data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}) @@ -571,7 +572,9 @@ async def test_update_service_account_requires_team_id(): @pytest.mark.asyncio async def test_update_service_account_works_with_team_id(): - data = UpdateKeyRequest(key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ") + data = UpdateKeyRequest( + key="sk-1", metadata={"service_account_id": "sa"}, team_id="IJ" + ) existing_key = LiteLLM_VerificationToken(token="hashed") await prepare_key_update_data(data=data, existing_key_row=existing_key) @@ -580,22 +583,22 @@ async def test_update_service_account_works_with_team_id(): @pytest.mark.asyncio async def test_validate_team_id_used_in_service_account_request_requires_team_id(): """ - Test that validate_team_id_used_in_service_account_request raises HTTPException + Test that validate_team_id_used_in_service_account_request raises HTTPException when team_id is None for service account key generation. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_team_id_used_in_service_account_request, ) - + mock_prisma_client = AsyncMock() - + # Test that HTTPException is raised when team_id is None with pytest.raises(HTTPException) as exc_info: await validate_team_id_used_in_service_account_request( team_id=None, prisma_client=mock_prisma_client, ) - + assert exc_info.value.status_code == 400 assert "team_id is required for service account keys" in str(exc_info.value.detail) @@ -603,7 +606,7 @@ async def test_validate_team_id_used_in_service_account_request_requires_team_id @pytest.mark.asyncio async def test_validate_team_id_used_in_service_account_request_requires_prisma_client(): """ - Test that validate_team_id_used_in_service_account_request raises HTTPException + Test that validate_team_id_used_in_service_account_request raises HTTPException when prisma_client is None for service account key generation. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( @@ -616,78 +619,76 @@ async def test_validate_team_id_used_in_service_account_request_requires_prisma_ team_id="test-team-id", prisma_client=None, ) - + assert exc_info.value.status_code == 400 - assert "prisma_client is required for service account keys" in str(exc_info.value.detail) + assert "prisma_client is required for service account keys" in str( + exc_info.value.detail + ) @pytest.mark.asyncio async def test_validate_team_id_used_in_service_account_request_checks_team_exists(): """ - Test that validate_team_id_used_in_service_account_request validates that + Test that validate_team_id_used_in_service_account_request validates that the team_id exists in the database for service account key generation. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_team_id_used_in_service_account_request, ) - + mock_prisma_client = AsyncMock() - + # Mock the database query to return None (team doesn't exist) mock_find_unique = AsyncMock(return_value=None) mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique - + # Test that HTTPException is raised when team doesn't exist in DB with pytest.raises(HTTPException) as exc_info: await validate_team_id_used_in_service_account_request( team_id="non-existent-team-id", prisma_client=mock_prisma_client, ) - + assert exc_info.value.status_code == 400 assert "team_id does not exist in the database" in str(exc_info.value.detail) - + # Verify the database was queried with the correct parameters - mock_find_unique.assert_called_once_with( - where={"team_id": "non-existent-team-id"} - ) + mock_find_unique.assert_called_once_with(where={"team_id": "non-existent-team-id"}) @pytest.mark.asyncio async def test_validate_team_id_used_in_service_account_request_success(): """ - Test that validate_team_id_used_in_service_account_request returns True + Test that validate_team_id_used_in_service_account_request returns True when team_id exists in the database for service account key generation. """ from litellm.proxy.management_endpoints.key_management_endpoints import ( validate_team_id_used_in_service_account_request, ) - + mock_prisma_client = AsyncMock() - + # Mock the database query to return a team object (team exists) mock_team = {"team_id": "existing-team-id", "team_name": "Test Team"} mock_find_unique = AsyncMock(return_value=mock_team) mock_prisma_client.db.litellm_teamtable.find_unique = mock_find_unique - + # Test that function returns True when team exists result = await validate_team_id_used_in_service_account_request( team_id="existing-team-id", prisma_client=mock_prisma_client, ) - + assert result is True - + # Verify the database was queried with the correct parameters - mock_find_unique.assert_called_once_with( - where={"team_id": "existing-team-id"} - ) + mock_find_unique.assert_called_once_with(where={"team_id": "existing-team-id"}) @pytest.mark.asyncio async def test_generate_service_account_key_endpoint_validation(): """ - Test that the /key/service-account/generate endpoint properly validates + Test that the /key/service-account/generate endpoint properly validates team_id requirement and team existence in database. """ from unittest.mock import patch @@ -705,16 +706,16 @@ async def test_generate_service_account_key_endpoint_validation(): ), litellm_changed_by=None, ) - + assert exc_info.value.status_code == 400 assert "team_id is required for service account keys" in str(exc_info.value.detail) - - # Test case 2: Team doesn't exist in database - with patch('litellm.proxy.proxy_server.prisma_client') as mock_prisma: + + # Test case 2: Team doesn't exist in database + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # Mock team not found mock_find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_teamtable.find_unique = mock_find_unique - + with pytest.raises(HTTPException) as exc_info: await generate_service_account_key_fn( data=GenerateKeyRequest(team_id="non-existent-team"), @@ -723,7 +724,165 @@ async def test_generate_service_account_key_endpoint_validation(): ), litellm_changed_by=None, ) - + assert exc_info.value.status_code == 400 assert "team_id does not exist in the database" in str(exc_info.value.detail) + +@pytest.mark.asyncio +async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch): + """ + Test that the unblock_key endpoint correctly handles both sk- prefixed tokens + and hashed tokens by properly converting sk- tokens to hashed format before + database operations. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key + + # Mock dependencies + mock_prisma_client = AsyncMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + + # Use a proper 64-character hex hash for testing + test_hashed_token = ( + "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + ) + + # Mock the key record that will be returned from database + mock_key_record = MagicMock() + mock_key_record.token = test_hashed_token + mock_key_record.blocked = False + mock_key_record.model_dump_json.return_value = ( + f'{{"token": "{test_hashed_token}", "blocked": false}}' + ) + + # Mock database operations + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=mock_key_record + ) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=mock_key_record + ) + + # Mock get_key_object and _cache_key_object functions + mock_key_object = MagicMock() + mock_key_object.blocked = True # Initially blocked + + # Mock hash_token function + def mock_hash_token(token): + if token == "sk-test123456789": + return test_hashed_token + return token + + # Apply monkeypatch + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache + ) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj + ) + monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", mock_hash_token) + monkeypatch.setattr( + "litellm.store_audit_logs", False + ) # Disable audit logs for simpler test + + # Mock get_key_object and _cache_key_object + async def mock_get_key_object(**kwargs): + return mock_key_object + + async def mock_cache_key_object(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints.get_key_object", + mock_get_key_object, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._cache_key_object", + mock_cache_key_object, + ) + + # Create mock request and user auth + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user" + ) + + # Test Case 1: Using sk- prefixed token + sk_token_request = BlockKeyRequest(key="sk-test123456789") + + result = await unblock_key( + data=sk_token_request, + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # Verify that the database update was called with hashed token + mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with( + where={"token": test_hashed_token}, data={"blocked": False} + ) + + assert result == mock_key_record + assert mock_key_object.blocked == False # Should be updated to unblocked + + # Reset mocks for second test + mock_prisma_client.db.litellm_verificationtoken.update.reset_mock() + mock_key_object.blocked = True # Reset to blocked state + + # Test Case 2: Using already hashed token + hashed_token_request = BlockKeyRequest(key=test_hashed_token) + + result = await unblock_key( + data=hashed_token_request, + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # Verify that the database update was called with the same hashed token + mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with( + where={"token": test_hashed_token}, data={"blocked": False} + ) + + assert result == mock_key_record + assert mock_key_object.blocked == False # Should be updated to unblocked + + +@pytest.mark.asyncio +async def test_unblock_key_invalid_key_format(monkeypatch): + """ + Test that unblock_key properly validates key format and raises appropriate errors + for invalid keys. + """ + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key + from litellm.proxy.utils import ProxyException + + # Mock prisma_client to avoid DB connection error + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Mock request and user auth + mock_request = MagicMock() + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user" + ) + + # Test with invalid key format + invalid_key_request = BlockKeyRequest(key="invalid-key-format") + + with pytest.raises(ProxyException) as exc_info: + await unblock_key( + data=invalid_key_request, + http_request=mock_request, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + assert exc_info.value.code == "400" + assert "Invalid key format" in str(exc_info.value.message)