diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a550af31e7..5e573d0bbe 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2324,8 +2324,6 @@ async def bulk_update_keys( error_message = error_detail.get("error", str(e)) else: error_message = str(error_detail) - elif isinstance(e, ProxyException): - error_message = e.message else: error_message = str(e) 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 37551fe387..48a1f7936c 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 @@ -5108,8 +5108,6 @@ async def test_get_and_validate_existing_key(): """ from fastapi import HTTPException - from litellm.proxy.utils import ProxyException - # Test Case 1: Successfully retrieve existing key mock_prisma_client = AsyncMock() mock_key = LiteLLM_VerificationToken( @@ -5118,37 +5116,31 @@ async def test_get_and_validate_existing_key(): models=["gpt-4"], team_id=None, ) - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=mock_key) - - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - return_value="hashed-test-key-123", - ): - result = await _get_and_validate_existing_key( - token="test-key-123", + mock_prisma_client.get_data = AsyncMock(return_value=mock_key) + + result = await _get_and_validate_existing_key( + token="test-key-123", + prisma_client=mock_prisma_client, + ) + + assert result == mock_key + mock_prisma_client.get_data.assert_called_once_with( + token="test-key-123", + table_name="key", + query_type="find_unique", + ) + + # Test Case 2: Key not found raises HTTPException + mock_prisma_client.get_data = AsyncMock(return_value=None) + + with pytest.raises(HTTPException) as exc_info: + await _get_and_validate_existing_key( + token="non-existent-key", prisma_client=mock_prisma_client, ) - - assert result == mock_key - mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with( - where={"token": "hashed-test-key-123"} - ) - - # Test Case 2: Key not found raises ProxyException - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) - - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - return_value="hashed-non-existent-key", - ): - with pytest.raises(ProxyException) as exc_info: - await _get_and_validate_existing_key( - token="non-existent-key", - prisma_client=mock_prisma_client, - ) - - assert str(exc_info.value.code) == "404" - assert "Key not found" in str(exc_info.value.message) + + assert exc_info.value.status_code == 404 + assert "Key not found" in str(exc_info.value.detail) # Test Case 3: Database not connected raises HTTPException with pytest.raises(HTTPException) as exc_info: @@ -5197,80 +5189,75 @@ async def test_process_single_key_update(): "tags": ["production"], } - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key) + mock_prisma_client.get_data = AsyncMock(return_value=existing_key) mock_updated_key_obj = MagicMock() mock_updated_key_obj.model_dump.return_value = updated_key_data mock_prisma_client.update_data = AsyncMock( return_value={"data": mock_updated_key_obj} ) - + # Mock prepare_key_update_data with patch( "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data" ) as mock_prepare: mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]} - + # Mock TeamMemberPermissionChecks with patch( "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint" ) as mock_permission_check: mock_permission_check.return_value = None - + # Mock _delete_cache_key_object with patch( "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object" ) as mock_delete_cache: mock_delete_cache.return_value = None - + # Mock hash_token (imported from litellm.proxy._types) with patch( "litellm.proxy._types.hash_token" ) as mock_hash: mock_hash.return_value = "hashed-test-key-123" - - # Mock _hash_token_if_needed + + # Mock KeyManagementEventHooks with patch( - "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - return_value="hashed-test-key-123", + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" ): - # Mock KeyManagementEventHooks - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" - ): - # Create update request - key_update_item = BulkUpdateKeyRequestItem( - key="test-key-123", - max_budget=100.0, - tags=["production"], - ) - - user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin-user", - ) - - # Call the function - result = await _process_single_key_update( - key_update_item=key_update_item, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - prisma_client=mock_prisma_client, - user_api_key_cache=mock_user_api_key_cache, - proxy_logging_obj=mock_proxy_logging_obj, - llm_router=mock_llm_router, - ) - - # Verify results - assert result is not None - assert "token" not in result # Token should be removed - assert result.get("max_budget") == 100.0 - assert result.get("tags") == ["production"] - - # Verify mocks were called - mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once() - mock_prisma_client.update_data.assert_called_once() - mock_delete_cache.assert_called_once() + # Create update request + key_update_item = BulkUpdateKeyRequestItem( + key="test-key-123", + max_budget=100.0, + tags=["production"], + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + # Call the function + result = await _process_single_key_update( + key_update_item=key_update_item, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + prisma_client=mock_prisma_client, + user_api_key_cache=mock_user_api_key_cache, + proxy_logging_obj=mock_proxy_logging_obj, + llm_router=mock_llm_router, + ) + + # Verify results + assert result is not None + assert "token" not in result # Token should be removed + assert result.get("max_budget") == 100.0 + assert result.get("tags") == ["production"] + + # Verify mocks were called + mock_prisma_client.get_data.assert_called_once() + mock_prisma_client.update_data.assert_called_once() + mock_delete_cache.assert_called_once() @pytest.mark.asyncio @@ -5460,17 +5447,15 @@ async def test_bulk_update_keys_partial_failures(monkeypatch): } # First key exists, second key doesn't exist - mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + mock_prisma_client.get_data = AsyncMock( side_effect=[existing_key_1, None] # Second key not found ) - # get_data is used by the error handler to fetch key info for failed updates - mock_prisma_client.get_data = AsyncMock(return_value=None) mock_updated_key_1_obj = MagicMock() mock_updated_key_1_obj.model_dump.return_value = updated_key_1_data mock_prisma_client.update_data = AsyncMock( return_value={"data": mock_updated_key_1_obj} ) - + # Patch dependencies monkeypatch.setattr( "litellm.proxy.proxy_server.prisma_client", mock_prisma_client @@ -5482,13 +5467,13 @@ async def test_bulk_update_keys_partial_failures(monkeypatch): "litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj ) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_llm_router) - + # Mock helper functions with patch( "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data" ) as mock_prepare: mock_prepare.return_value = {"max_budget": 100.0, "tags": ["production"]} - + with patch( "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint" ): @@ -5499,50 +5484,46 @@ async def test_bulk_update_keys_partial_failures(monkeypatch): "litellm.proxy._types.hash_token" ) as mock_hash: mock_hash.return_value = "hashed-key-1" - + with patch( - "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - side_effect=lambda token: f"hashed-{token}", + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" ): - with patch( - "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" - ): - # Create request with one valid and one invalid key - request_data = BulkUpdateKeyRequest( - keys=[ - BulkUpdateKeyRequestItem( - key="test-key-1", - max_budget=100.0, - tags=["production"], - ), - BulkUpdateKeyRequestItem( - key="non-existent-key", - max_budget=200.0, - tags=["staging"], - ), - ] - ) - - user_api_key_dict = UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - api_key="sk-admin", - user_id="admin-user", - ) - - # Call endpoint - response = await bulk_update_keys( - data=request_data, - user_api_key_dict=user_api_key_dict, - litellm_changed_by=None, - ) - - # Verify response - assert response.total_requested == 2 - assert len(response.successful_updates) == 1 - assert len(response.failed_updates) == 1 - assert response.successful_updates[0].key == "test-key-1" - assert response.failed_updates[0].key == "non-existent-key" - assert "Key not found" in response.failed_updates[0].failed_reason + # Create request with one valid and one invalid key + request_data = BulkUpdateKeyRequest( + keys=[ + BulkUpdateKeyRequestItem( + key="test-key-1", + max_budget=100.0, + tags=["production"], + ), + BulkUpdateKeyRequestItem( + key="non-existent-key", + max_budget=200.0, + tags=["staging"], + ), + ] + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + # Call endpoint + response = await bulk_update_keys( + data=request_data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + ) + + # Verify response + assert response.total_requested == 2 + assert len(response.successful_updates) == 1 + assert len(response.failed_updates) == 1 + assert response.successful_updates[0].key == "test-key-1" + assert response.failed_updates[0].key == "non-existent-key" + assert "Key not found" in response.failed_updates[0].failed_reason @pytest.mark.parametrize(