diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index ea19befccc6..1c4eada67fb 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -217,7 +217,7 @@ async def handle_budget_for_entity( ) # Get all budget field names - budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() + budget_params: Final = LiteLLM_BudgetTable.model_fields.keys() - {"budget_id"} # Extract budget fields from data _json_data: Final = _BUDGET_DATA_MAPPING.validate_python( diff --git a/tests/integration/management/test_tag_budget_update.py b/tests/integration/management/test_tag_budget_update.py new file mode 100644 index 00000000000..c8ce1bdc8f4 --- /dev/null +++ b/tests/integration/management/test_tag_budget_update.py @@ -0,0 +1,20 @@ +from typing import Final +from uuid import uuid4 + +from tests.integration._support.client import Gateway, object_value, string_value + + +def test_tag_budget_can_be_updated_with_its_existing_budget_id(gateway: Gateway) -> None: + name: Final = f"integration-tag-{uuid4().hex}" + gateway.post("/tag/new", {"name": name, "max_budget": 10}) + with gateway.scenario() as scenario: + scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": name}) + before: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name]) + budget_id: Final = string_value(object_value(before["litellm_budget_table"])["budget_id"]) + scenario.cleanups.callback(gateway.post, "/budget/delete", {"id": budget_id}) + response: Final = gateway.request( + "POST", "/tag/update", {"name": name, "budget_id": budget_id, "max_budget": 20} + ) + assert response.status_code == 200, response.text + after: Final = object_value(gateway.post("/tag/info", {"names": [name]})[name]) + assert object_value(after["litellm_budget_table"])["max_budget"] == 20, after diff --git a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py index 644f0656c7a..640e17aa58e 100644 --- a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py @@ -249,6 +249,35 @@ async def test_add_new_member_links_default_team_budget_id(): assert create_data["budget_id"] == test_default_budget_id +@pytest.mark.asyncio +async def test_handle_budget_for_entity_updates_with_existing_budget_id_in_request() -> None: + from litellm.proxy.management_helpers.utils import handle_budget_for_entity + from litellm.types.tag_management import TagUpdateRequest + + budget_update: Final = AsyncMock() + prisma_client: Final = cast( + PrismaClient, + SimpleNamespace(db=SimpleNamespace(litellm_budgettable=SimpleNamespace(update=budget_update))), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), + ): + budget_id: Final = await handle_budget_for_entity( + data=TagUpdateRequest(name="budget-tag", budget_id="budget-1", max_budget=20), + existing_budget_id="budget-1", + user_api_key_dict=UserAPIKeyAuth(user_id="admin"), + prisma_client=prisma_client, + litellm_proxy_admin_name="admin", + ) + + assert budget_id == "budget-1" + budget_update.assert_awaited_once_with( + where={"budget_id": "budget-1"}, + data={"budget_id": "budget-1", "max_budget": 20.0, "updated_by": "admin"}, + ) + + @pytest.mark.parametrize( ("request_fields", "cleared_budget_fields", "should_update", "expected_max_budget"), cast(