test: assert the budget-authority rejection on owner regenerate (test-quality)

Generated-by: termag

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,

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