From ca353b4698a1aeb64f5a8ef520e9ac5b957a2c72 Mon Sep 17 00:00:00 2001 From: Souravrajvi0 Date: Tue, 21 Jul 2026 15:03:18 +0000 Subject: [PATCH] test(proxy): cover existing-budget duration update on customer update Strengthen Greptile P2 coverage for #33941: assert budget table update where clause, and add regression for max_budget + budget_duration on an existing linked budget. --- .../test_customer_budget.py | 57 ++++++++++++++++++- 1 file changed, 56 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py b/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py index 410efe6475f..d7b27895910 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_budget.py @@ -500,12 +500,67 @@ async def test_update_customer_updates_existing_budget_with_duration_and_reset_a await update_end_user(update_request, mock_user_api_key_dict) mock_prisma_client.db.litellm_budgettable.update.assert_called_once() - update_data = mock_prisma_client.db.litellm_budgettable.update.call_args[1]["data"] + update_call = mock_prisma_client.db.litellm_budgettable.update.call_args + assert update_call[1]["where"] == {"budget_id": "existing-budget-123"} + update_data = update_call[1]["data"] assert update_data["budget_duration"] == "1mo" assert update_data["budget_reset_at"] is not None assert not mock_prisma_client.db.litellm_budgettable.create.called +@pytest.mark.asyncio +@patch("litellm.proxy.proxy_server.prisma_client") +@patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin") +async def test_update_customer_updates_existing_budget_max_budget_and_duration( + mock_prisma_client, mock_user_api_key_dict, mock_existing_customer +): + """ + Regression for #33941 / Greptile P2: updating an existing linked budget via + /customer/update must persist both max_budget and budget_duration together. + """ + mock_existing_customer.model_dump.return_value = { + "user_id": "test-bug-repro@example.com", + "blocked": False, + "litellm_budget_table": { + "budget_id": "existing-budget-123", + "max_budget": 25.0, + "budget_duration": "7d", + }, + } + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=mock_existing_customer + ) + mock_prisma_client.db.litellm_budgettable.update = AsyncMock( + return_value=MagicMock(budget_id="existing-budget-123") + ) + + mock_updated_user = MagicMock() + mock_updated_user.model_dump.return_value = { + "user_id": "test-bug-repro@example.com", + "blocked": False, + "budget_id": "existing-budget-123", + } + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + return_value=mock_updated_user + ) + + update_request = UpdateCustomerRequest( + user_id="test-bug-repro@example.com", + max_budget=100.0, + budget_duration="30d", + ) + + await update_end_user(update_request, mock_user_api_key_dict) + + mock_prisma_client.db.litellm_budgettable.update.assert_called_once() + update_data = mock_prisma_client.db.litellm_budgettable.update.call_args[1]["data"] + assert update_data["max_budget"] == 100.0 + assert update_data["budget_duration"] == "30d" + assert update_data["budget_reset_at"] is not None + assert not mock_prisma_client.db.litellm_budgettable.create.called + + @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.prisma_client") @patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")