From 3ef6f96e96c67b571ed647ca469f10a9f300dfd7 Mon Sep 17 00:00:00 2001 From: Jack Danger Date: Sun, 18 Oct 2026 09:29:00 +0000 Subject: [PATCH] test: assert the budget-authority rejection on owner regenerate (test-quality) Generated-by: termag --- .../key_management_endpoints.py | 29 ++++------- .../test_key_management_endpoints.py | 48 +++++++++---------- 2 files changed, 33 insertions(+), 44 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index fc7f45d7dd4..5a234400af1 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -934,7 +934,7 @@ def _strip_non_admin_changed_by_metadata( audit rows, so a self-granted flag lets any key holder forge the primary audit attribution field. Silently stripped for non-admin writes. """ - metadata = data.metadata + metadata: Final = data.metadata if not isinstance(metadata, dict): return if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: @@ -3342,9 +3342,7 @@ async def _validate_end_user_budget_id_change( raise HTTPException(status_code=400, detail=missing_detail) -_TEMP_BUDGET_METADATA_KEYS: Final[frozenset[str]] = frozenset( - {"temp_budget_increase", "temp_budget_expiry"} -) +_TEMP_BUDGET_METADATA_KEYS: Final[frozenset[str]] = frozenset({"temp_budget_increase", "temp_budget_expiry"}) def _key_request_changes_budget( @@ -3359,13 +3357,11 @@ def _key_request_changes_budget( in. /key/update and /key/regenerate treat all of these as admin-only changes for non-admin callers. """ - changed_max_budget: Final = ( - data.max_budget is not None - and data.max_budget != getattr(existing_key_row, "max_budget", None) + changed_max_budget: Final = data.max_budget is not None and data.max_budget != getattr( + existing_key_row, "max_budget", None ) - changed_budget_duration: Final = ( - data.budget_duration is not None - and data.budget_duration != getattr(existing_key_row, "budget_duration", None) + changed_budget_duration: Final = data.budget_duration is not None and data.budget_duration != getattr( + existing_key_row, "budget_duration", None ) return ( changed_max_budget @@ -3565,9 +3561,7 @@ async def _validate_update_key_data( # through the non-admin path would let a key owner / team member # overwrite the live counter below real usage and silently weaken # enforcement. - _is_budget_change: Final = _key_request_changes_budget( - data=data, existing_key_row=existing_key_row - ) + _is_budget_change: Final = _key_request_changes_budget(data=data, existing_key_row=existing_key_row) _existing_metadata: Final = getattr(existing_key_row, "metadata", None) _existing_throttle: Final = ( @@ -6036,9 +6030,7 @@ async def _execute_virtual_key_regeneration( user_api_key_dict.user_role is not None and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value ) - if not _regenerate_is_proxy_admin and _key_request_changes_budget( - data=data, existing_key_row=key_in_db - ): + if not _regenerate_is_proxy_admin and _key_request_changes_budget(data=data, existing_key_row=key_in_db): await _check_key_admin_access( user_api_key_dict=user_api_key_dict, hashed_token=hashed_api_key, @@ -8157,10 +8149,7 @@ async def _can_user_query_key_info( or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value ) or user_api_key_dict.api_key == key - or ( - key_info.user_id is not None - and key_info.user_id == user_api_key_dict.user_id - ) + or (key_info.user_id is not None and key_info.user_id == user_api_key_dict.user_id) or await TeamMemberPermissionChecks.user_belongs_to_keys_team( user_api_key_dict=user_api_key_dict, existing_key_row=key_info, diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 4fb14ce19c0..a82891d4680 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -23247,12 +23247,6 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): ) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _make_regenerate_mock_prisma()) - admin_check = AsyncMock(return_value=None) - monkeypatch.setattr( - "litellm.proxy.management_endpoints.key_management_endpoints._check_key_admin_access", - admin_check, - ) - existing_key = _make_regenerate_existing_key() user_api_key_cache = UserApiKeyCache() caller = UserAPIKeyAuth( @@ -23261,6 +23255,14 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): user_id="user-1", ) + # The real _check_key_admin_access runs: its lookup resolves the target key + # through the mocked prisma client, its team step is skipped because the key + # has no team, and the internal user is told the route is admin-only. + prisma_client = _make_regenerate_mock_prisma() + prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=_make_regenerate_existing_key() + ) + with ( patch( "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", @@ -23279,25 +23281,23 @@ async def test_key_regeneration_budget_change_requires_admin(monkeypatch): "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook" ), ): - await _execute_virtual_key_regeneration( - prisma_client=_make_regenerate_mock_prisma(), - key_in_db=existing_key, - hashed_api_key="abc123", - key="abc123", - data=RegenerateKeyRequest(spend=0.0, max_budget=100000.0), - user_api_key_dict=caller, - litellm_changed_by=None, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=AsyncMock(), - ) + with pytest.raises(HTTPException) as exc: + await _execute_virtual_key_regeneration( + prisma_client=prisma_client, + key_in_db=existing_key, + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(spend=0.0, max_budget=100000.0), + user_api_key_dict=caller, + litellm_changed_by=None, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=AsyncMock(), + ) - assert admin_check.await_count == 1 - assert ( - "/key/regenerate (max_budget/spend)" in admin_check.call_args.kwargs["route"] - ) - # the caller also observes that the raise carried the guarded route label - # through _check_key_admin_access's HTTPException construction - assert admin_check.call_args.kwargs.get("prisma_client") is not None + # The caller observes the rejection itself: a budget-moving regenerate for a + # non-admin cannot silently reset spend or lift max_budget. + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "/key/regenerate (max_budget/spend)" in str(exc.value.detail) @pytest.mark.asyncio