From 35c00b7093897f227e7ccfccf25fc2ef0dcddad8 Mon Sep 17 00:00:00 2001 From: jibanez-staticduo Date: Tue, 22 Sep 2026 13:55:22 +0200 Subject: [PATCH] test(proxy): stop the proxy budget from leaking between proxy tests `proxy-infra / Run tests` failed on four parametrizations of `test_real_proxy_child_auth_privacy_and_body_policy` -- the four that expect the child request to succeed. The traceback ends in upstream code this PR does not touch: `_user_api_key_auth_builder` calls `_fetch_global_spend_with_event_coordination`, whose loader reads `prisma_client.db.litellm_usertable`, and the fixture installs `object()` as the prisma client, so the read raises `'object' object has no attribute 'db'` and authentication returns 401. That block only runs when `litellm.max_budget > 0`, and the fixture never sets it, so the value came from a test that ran earlier on the same xdist worker. `test_add_proxy_budget_to_db_only_creates_user_no_keys` and `test_add_proxy_budget_to_db_backfills_budget_reset_at` assign `litellm.max_budget = 100.0` and `litellm.budget_duration` directly and never restore them, and which tests share a worker changes from run to run, which is why the same commit was green for these four cases and red for them later. Two small test fixes, no production change: - the two budget startup tests now set those globals through `monkeypatch`, so they are undone at teardown like the rest of the file's budget tests already do; - the compaction fixture pins `litellm.max_budget` to 0.0 beside the other pins, so it asserts its own precondition instead of depending on suite order. Validation: `test_native_compaction.py` plus both budget tests -- 16 passed; `ruff check --config ruff-tests.toml` clean on both files, and `ruff format --diff` leaves the added lines alone (the format check already reported this file before the change). --- .../proxy/test_native_compaction.py | 6 ++++++ tests/test_litellm/proxy/test_proxy_server.py | 17 ++++++++++------- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/proxy/test_native_compaction.py b/tests/test_litellm/proxy/test_native_compaction.py index d24624aacc5..b8062c6a4de 100644 --- a/tests/test_litellm/proxy/test_native_compaction.py +++ b/tests/test_litellm/proxy/test_native_compaction.py @@ -7,6 +7,7 @@ import pytest from fastapi import FastAPI, Request from pydantic import TypeAdapter +import litellm from litellm.caching.caching import DualCache from litellm.exceptions import BadRequestError from litellm.litellm_core_utils.initialize_dynamic_callback_params import inherit_message_logging_privacy @@ -118,6 +119,11 @@ async def test_real_proxy_child_auth_privacy_and_body_policy( monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(proxy_server, "general_settings", {}) + # Pin the proxy-wide budget: authentication only reads the global spend when a + # proxy max budget is configured, and that read goes through the stub prisma + # client above. A budget left set by an earlier test on the same worker would + # turn this fixture's child requests into 401s. + monkeypatch.setattr(litellm, "max_budget", 0.0) monkeypatch.setattr(common_request_processing, "route_request", route) with inherit_message_logging_privacy(True): call: Final = with_proxy_compaction_executor( diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 84189cd3305..6d43870819d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3200,7 +3200,7 @@ def test_normalize_datetime_for_sorting(): @pytest.mark.asyncio -async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): +async def test_add_proxy_budget_to_db_only_creates_user_no_keys(monkeypatch: pytest.MonkeyPatch): """ Test that _add_proxy_budget_to_db only creates a user and no keys are added. @@ -3217,9 +3217,12 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): import litellm from litellm.proxy.proxy_server import ProxyStartupEvent - # Set up required litellm settings - litellm.budget_duration = "30d" - litellm.max_budget = 100.0 + # Set up required litellm settings. Through monkeypatch rather than plain + # assignment: `litellm.max_budget` is process-global, and any later test on + # this worker that authenticates reads the global proxy spend whenever a + # proxy budget is set, which needs a real prisma client. + monkeypatch.setattr(litellm, "budget_duration", "30d") + monkeypatch.setattr(litellm, "max_budget", 100.0) litellm_proxy_budget_name = "litellm-proxy-budget" @@ -3258,7 +3261,7 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): @pytest.mark.asyncio -async def test_add_proxy_budget_to_db_backfills_budget_reset_at(): +async def test_add_proxy_budget_to_db_backfills_budget_reset_at(monkeypatch: pytest.MonkeyPatch): """ Test that _upsert_proxy_budget_with_reset_at_backfill issues a conditional update_many with `WHERE budget_reset_at IS NULL` to backfill the column on @@ -3276,8 +3279,8 @@ async def test_add_proxy_budget_to_db_backfills_budget_reset_at(): import litellm from litellm.proxy.proxy_server import ProxyStartupEvent - litellm.budget_duration = "30d" - litellm.max_budget = 100.0 + monkeypatch.setattr(litellm, "budget_duration", "30d") + monkeypatch.setattr(litellm, "max_budget", 100.0) litellm_proxy_budget_name = "litellm-proxy-budget" mock_prisma = MagicMock()