From 1181f1bb2e54df05be79802084ec6d567e83a236 Mon Sep 17 00:00:00 2001 From: Jason Jensen Date: Wed, 30 Jul 2025 16:47:25 +0000 Subject: [PATCH] test(customer): add unit tests for budget_id handling in customer update endpoint --- .../test_customer_update_budget_fix.py | 340 ++++++++++++++++++ 1 file changed, 340 insertions(+) create mode 100644 tests/test_litellm/test_customer_update_budget_fix.py diff --git a/tests/test_litellm/test_customer_update_budget_fix.py b/tests/test_litellm/test_customer_update_budget_fix.py new file mode 100644 index 0000000000..96ce24d48c --- /dev/null +++ b/tests/test_litellm/test_customer_update_budget_fix.py @@ -0,0 +1,340 @@ +""" +Unit tests for customer update budget fix. + +Tests the specific changes made to fix the customer update endpoint's budget_id handling: +1. budget_id routing to end user table data (not budget table data) +2. Correct Prisma include field usage ("end_users" not "litellm_endusertable") +3. Required fields (created_by, updated_by) in budget creation +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_EndUserTable, + UpdateCustomerRequest, +) +from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.proxy.management_endpoints.customer_endpoints import update_end_user + + +class TestCustomerUpdateBudgetFix: + """Test class for customer update budget fix functionality.""" + + @pytest.fixture + def mock_user_api_key_dict(self): + """Mock user API key auth object.""" + mock_auth = MagicMock(spec=UserAPIKeyAuth) + mock_auth.user_id = "test-admin-user" + return mock_auth + + @pytest.fixture + def mock_existing_customer(self): + """Mock existing customer data.""" + return MagicMock(spec=LiteLLM_EndUserTable) + + @pytest.fixture + def mock_budget_table(self): + """Mock budget table data.""" + return MagicMock(spec=LiteLLM_BudgetTable) + + @pytest.mark.asyncio + @patch('litellm.proxy.proxy_server.prisma_client') + @patch('litellm.proxy.proxy_server.litellm_proxy_admin_name', 'admin') + async def test_budget_id_routes_to_end_user_data_not_budget_data( + self, + mock_prisma_client, + mock_user_api_key_dict, + mock_existing_customer + ): + """ + Test that budget_id is correctly routed to end user table data, not budget table data. + + This was the main bug - budget_id was being treated as budget creation data + instead of budget linking data. + """ + # Arrange + mock_existing_customer.model_dump.return_value = { + "user_id": "test-user", + "blocked": False, + "litellm_budget_table": None + } + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=mock_existing_customer + ) + + mock_updated_user = MagicMock() + mock_updated_user.model_dump.return_value = { + "user_id": "test-user", + "budget_id": "existing-budget-123", + "blocked": False + } + + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + return_value=mock_updated_user + ) + + # Create update request with only budget_id (no other budget fields) + update_request = UpdateCustomerRequest( + user_id="test-user", + budget_id="existing-budget-123" + ) + + # Act + result = await update_end_user(update_request, mock_user_api_key_dict) + + # Assert + # Verify that update was called on end user table with budget_id + mock_prisma_client.db.litellm_endusertable.update.assert_called_once() + call_args = mock_prisma_client.db.litellm_endusertable.update.call_args + + # Check that budget_id is in the update data for end user table + update_data = call_args[1]['data'] # kwargs['data'] + assert 'budget_id' in update_data + assert update_data['budget_id'] == "existing-budget-123" + + # Verify that NO budget creation was attempted + 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_budget_creation_uses_correct_include_field( + self, + mock_prisma_client, + mock_user_api_key_dict, + mock_existing_customer + ): + """ + Test that budget creation uses correct Prisma include field "end_users". + + The bug was using "litellm_endusertable" instead of "end_users". + """ + # Arrange + mock_existing_customer.model_dump.return_value = { + "user_id": "test-user", + "blocked": False, + "litellm_budget_table": None + } + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=mock_existing_customer + ) + + # Mock budget creation + mock_created_budget = MagicMock() + mock_created_budget.budget_id = "new-budget-456" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_created_budget + ) + + # Mock end user update + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + return_value=MagicMock() + ) + + # Create update request with budget creation fields (not just budget_id) + update_request = UpdateCustomerRequest( + user_id="test-user", + max_budget=100.0, # This triggers budget creation + rpm_limit=200 # Use valid budget field + ) + + # Act + await update_end_user(update_request, mock_user_api_key_dict) + + # Assert + # Verify budget creation was called with correct include field + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + call_args = mock_prisma_client.db.litellm_budgettable.create.call_args + + # Check that include uses correct field name "end_users" + include_param = call_args[1]['include'] # kwargs['include'] + assert 'end_users' in include_param + assert include_param['end_users'] is True + + # Ensure the old incorrect field is NOT used + assert 'litellm_endusertable' not in include_param + + @pytest.mark.asyncio + @patch('litellm.proxy.proxy_server.prisma_client') + @patch('litellm.proxy.proxy_server.litellm_proxy_admin_name', 'admin') + async def test_budget_creation_includes_required_fields( + self, + mock_prisma_client, + mock_user_api_key_dict, + mock_existing_customer + ): + """ + Test that budget creation includes required created_by and updated_by fields. + + These fields were missing and causing creation failures. + """ + # Arrange + mock_existing_customer.model_dump.return_value = { + "user_id": "test-user", + "blocked": False, + "litellm_budget_table": None + } + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=mock_existing_customer + ) + + # Mock budget creation + mock_created_budget = MagicMock() + mock_created_budget.budget_id = "new-budget-789" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_created_budget + ) + + # Mock end user update + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + return_value=MagicMock() + ) + + # Create update request with budget creation fields + update_request = UpdateCustomerRequest( + user_id="test-user", + max_budget=200.0 + ) + + # Act + await update_end_user(update_request, mock_user_api_key_dict) + + # Assert + # Verify budget creation was called with required fields + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + call_args = mock_prisma_client.db.litellm_budgettable.create.call_args + + # Check that created_by and updated_by are present in creation data + creation_data = call_args[1]['data'] # kwargs['data'] + assert 'created_by' in creation_data + assert 'updated_by' in creation_data + + # Verify the values are set correctly + assert creation_data['created_by'] == "test-admin-user" + assert creation_data['updated_by'] == "test-admin-user" + + # Verify budget fields are also included + assert 'max_budget' in creation_data + assert creation_data['max_budget'] == 200.0 + + @pytest.mark.asyncio + @patch('litellm.proxy.proxy_server.prisma_client') + @patch('litellm.proxy.proxy_server.litellm_proxy_admin_name', 'admin') + async def test_budget_creation_fallback_to_admin_name( + self, + mock_prisma_client, + mock_existing_customer + ): + """ + Test that budget creation falls back to litellm_proxy_admin_name when user_id is None. + """ + # Arrange - user with None user_id + mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth) + mock_user_api_key_dict.user_id = None + + mock_existing_customer.model_dump.return_value = { + "user_id": "test-user", + "blocked": False, + "litellm_budget_table": None + } + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=mock_existing_customer + ) + + # Mock budget creation + mock_created_budget = MagicMock() + mock_created_budget.budget_id = "new-budget-fallback" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_created_budget + ) + + # Mock end user update + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + return_value=MagicMock() + ) + + # Create update request with budget creation fields + update_request = UpdateCustomerRequest( + user_id="test-user", + max_budget=150.0, + tpm_limit=1000 # Add another budget field + ) + + # Act + await update_end_user(update_request, mock_user_api_key_dict) + + # Assert + # Verify budget creation was called with fallback admin name + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + call_args = mock_prisma_client.db.litellm_budgettable.create.call_args + + creation_data = call_args[1]['data'] # kwargs['data'] + assert creation_data['created_by'] == "admin" # litellm_proxy_admin_name + assert creation_data['updated_by'] == "admin" + + @pytest.mark.asyncio + @patch('litellm.proxy.proxy_server.prisma_client') + @patch('litellm.proxy.proxy_server.litellm_proxy_admin_name', 'admin') + async def test_budget_id_and_budget_fields_combination( + self, + mock_prisma_client, + mock_user_api_key_dict, + mock_existing_customer + ): + """ + Test behavior when both budget_id and budget creation fields are provided. + + budget_id should go to end user data, budget fields should trigger budget creation. + """ + # Arrange + mock_existing_customer.model_dump.return_value = { + "user_id": "test-user", + "blocked": False, + "litellm_budget_table": None + } + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock( + return_value=mock_existing_customer + ) + + # Mock budget creation + mock_created_budget = MagicMock() + mock_created_budget.budget_id = "new-budget-combo" + mock_prisma_client.db.litellm_budgettable.create = AsyncMock( + return_value=mock_created_budget + ) + + # Mock end user update + mock_updated_user = MagicMock() + mock_prisma_client.db.litellm_endusertable.update = AsyncMock( + return_value=mock_updated_user + ) + + # Create update request with both budget_id and budget creation fields + update_request = UpdateCustomerRequest( + user_id="test-user", + budget_id="existing-budget-link", # For linking to existing budget + max_budget=300.0, # This should trigger new budget creation + rpm_limit=500 # Use valid budget field + ) + + # Act + await update_end_user(update_request, mock_user_api_key_dict) + + # Assert + # Verify budget creation occurred (because max_budget was provided) + mock_prisma_client.db.litellm_budgettable.create.assert_called_once() + + # Verify end user update was called + mock_prisma_client.db.litellm_endusertable.update.assert_called_once() + call_args = mock_prisma_client.db.litellm_endusertable.update.call_args + + # The update data should contain budget_id from the created budget, not the original budget_id + update_data = call_args[1]['data'] + assert update_data['budget_id'] == "new-budget-combo" # From created budget \ No newline at end of file