diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 22dd3b6fca..fb600cee26 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -46,13 +46,35 @@ class ResetBudgetJob: await self.reset_budget_for_litellm_teams() ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ### - await self.reset_budget_for_litellm_endusers() + await self.reset_budget_for_litellm_budget_table() - async def reset_budget_for_litellm_endusers(self): + async def reset_budget_for_litellm_team_members( + self, budgets_to_reset: List[LiteLLM_BudgetTableFull] + ): """ - Resets the budget for all LiteLLM End-Users (Customers) if their budget has expired + Resets the budget for all LiteLLM Team Members if their budget has expired + """ + return await self.prisma_client.db.litellm_teammembership.update_many( + where={ + "budget_id": { + "in": [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + } + }, + data={ + "spend": 0, + }, + ) + + async def reset_budget_for_litellm_budget_table(self): + """ + Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired The corresponding Budget duration is also updated. """ + now = datetime.now(timezone.utc) start_time = time.time() endusers_to_reset: Optional[List[LiteLLM_EndUserTable]] = None @@ -69,6 +91,7 @@ class ResetBudgetJob: budget = await ResetBudgetJob._reset_budget_reset_at_date( budget, now ) + await self.prisma_client.update_data( query_type="update_many", data_list=budgets_to_reset, @@ -78,7 +101,15 @@ class ResetBudgetJob: endusers_to_reset = await self.prisma_client.get_data( table_name="enduser", query_type="find_all", - budget_id_list=[budget.budget_id for budget in budgets_to_reset], + budget_id_list=[ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ], + ) + + await self.reset_budget_for_litellm_team_members( + budgets_to_reset=budgets_to_reset ) if endusers_to_reset is not None and len(endusers_to_reset) > 0: @@ -125,19 +156,19 @@ class ResetBudgetJob: self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, duration=end_time - start_time, - call_type="reset_budget_endusers", + call_type="reset_budget_budget_table", start_time=start_time, end_time=end_time, event_metadata={ - "num_budgets_found": len(budgets_to_reset) - if budgets_to_reset - else 0, + "num_budgets_found": ( + len(budgets_to_reset) if budgets_to_reset else 0 + ), "budgets_found": json.dumps( budgets_to_reset, indent=4, default=str ), - "num_endusers_found": len(endusers_to_reset) - if endusers_to_reset - else 0, + "num_endusers_found": ( + len(endusers_to_reset) if endusers_to_reset else 0 + ), "endusers_found": json.dumps( endusers_to_reset, indent=4, default=str ), @@ -163,15 +194,15 @@ class ResetBudgetJob: start_time=start_time, end_time=end_time, event_metadata={ - "num_budgets_found": len(budgets_to_reset) - if budgets_to_reset - else 0, + "num_budgets_found": ( + len(budgets_to_reset) if budgets_to_reset else 0 + ), "budgets_found": json.dumps( budgets_to_reset, indent=4, default=str ), - "num_endusers_found": len(endusers_to_reset) - if endusers_to_reset - else 0, + "num_endusers_found": ( + len(endusers_to_reset) if endusers_to_reset else 0 + ), "endusers_found": json.dumps( endusers_to_reset, indent=4, default=str ), @@ -465,7 +496,10 @@ class ResetBudgetJob: item.spend = 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: # Get standardized reset time based on budget duration - from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + from litellm.proxy.common_utils.timezone_utils import ( + get_budget_reset_time, + ) + item.budget_reset_at = get_budget_reset_time( budget_duration=item.budget_duration ) @@ -510,10 +544,13 @@ class ResetBudgetJob: @staticmethod async def _reset_budget_reset_at_date( budget: LiteLLM_BudgetTableFull, current_time: datetime - ) -> Optional[LiteLLM_BudgetTableFull]: + ) -> LiteLLM_BudgetTableFull: try: if budget.budget_duration is not None: - from litellm.litellm_core_utils.duration_parser import duration_in_seconds + from litellm.litellm_core_utils.duration_parser import ( + duration_in_seconds, + ) + duration_s = duration_in_seconds(duration=budget.budget_duration) # Fallback for existing budgets that do not have a budget_reset_at date set, ensuring the duration is taken into account diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index d12816bc33..16959f88f0 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -282,7 +282,15 @@ async def _update_litellm_setting( in_memory_var: The in-memory variable to update success_message: Message to return on success """ - from litellm.proxy.proxy_server import proxy_config + from litellm.proxy.proxy_server import proxy_config, store_model_in_db + + if store_model_in_db is not True: + raise HTTPException( + status_code=500, + detail={ + "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + }, + ) # Update the in-memory settings in_memory_var = settings.model_dump(exclude_none=True) diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index a3c9523941..7cddde3042 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -243,10 +243,19 @@ async def test_reset_budget_endusers_partial_failure(): enduser["spend"] = 0.0 return enduser + async def fake_reset_team_members(budgets_to_reset): + return 1 + with patch.object( - ResetBudgetJob, "_reset_budget_for_enduser", side_effect=fake_reset_enduser - ) as mock_reset_enduser: - await job.reset_budget_for_litellm_endusers() + ResetBudgetJob, + "_reset_budget_for_enduser", + side_effect=fake_reset_enduser, + ) as mock_reset_enduser, patch.object( + ResetBudgetJob, + "reset_budget_for_litellm_team_members", + side_effect=fake_reset_team_members, + ) as mock_reset_team_members: + await job.reset_budget_for_litellm_budget_table() await asyncio.sleep(0.1) assert mock_reset_enduser.call_count == 6 @@ -415,6 +424,9 @@ async def test_reset_budget_continues_other_categories_on_failure(): enduser["spend"] = 0.0 return enduser + async def fake_reset_team_members(budgets_to_reset): + return 1 + with patch.object( ResetBudgetJob, "_reset_budget_for_key", side_effect=fake_reset_key ) as mock_reset_key, patch.object( @@ -423,7 +435,11 @@ async def test_reset_budget_continues_other_categories_on_failure(): ResetBudgetJob, "_reset_budget_for_team", side_effect=fake_reset_team ) as mock_reset_team, patch.object( ResetBudgetJob, "_reset_budget_for_enduser", side_effect=fake_reset_enduser - ) as mock_reset_enduser: + ) as mock_reset_enduser, patch.object( + ResetBudgetJob, + "reset_budget_for_litellm_team_members", + side_effect=fake_reset_team_members, + ) as mock_reset_team_members: # Call the overall reset_budget method. await job.reset_budget() await asyncio.sleep(0.1) @@ -432,7 +448,16 @@ async def test_reset_budget_continues_other_categories_on_failure(): called_tables = { call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list } - assert called_tables == {"key", "user", "team", "budget", "enduser"} + if mock_reset_team_members.call_count > 0: + called_tables.add("team_membership") + assert called_tables == { + "key", + "user", + "team", + "budget", + "enduser", + "team_membership", + } # Verify that update_data was called three times (one per category, enduser update includes two) assert prisma_client.update_data.await_count == 5 @@ -850,15 +875,22 @@ async def test_service_logger_endusers_success(): enduser["spend"] = 0.0 return enduser + async def fake_reset_team_members(budgets_to_reset): + return 1 + with patch.object( ResetBudgetJob, "_reset_budget_for_enduser", side_effect=fake_reset_enduser, - ): + ) as mock_reset_enduser, patch.object( + ResetBudgetJob, + "reset_budget_for_litellm_team_members", + side_effect=fake_reset_team_members, + ) as mock_reset_team_members: with patch( "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception" ) as mock_verbose_exc: - await job.reset_budget_for_litellm_endusers() + await job.reset_budget_for_litellm_budget_table() await asyncio.sleep(0.1) mock_verbose_exc.assert_not_called() @@ -876,7 +908,7 @@ async def test_service_logger_endusers_success(): @pytest.mark.asyncio -async def test_service_logger_users_failure(): +async def test_service_logger_endusers_failure(): """ Test that a failure during enduser reset calls the failure hook with appropriate metadata, logs the exception, and does not call the success hook. @@ -920,15 +952,22 @@ async def test_service_logger_users_failure(): enduser["spend"] = 0.0 return enduser + async def fake_reset_team_members(budgets_to_reset): + return 1 + with patch.object( ResetBudgetJob, "_reset_budget_for_enduser", side_effect=fake_reset_enduser, - ): + ) as mock_reset_enduser, patch.object( + ResetBudgetJob, + "reset_budget_for_litellm_team_members", + side_effect=fake_reset_team_members, + ) as mock_reset_team_members: with patch( "litellm.proxy.common_utils.reset_budget_job.verbose_proxy_logger.exception" ) as mock_verbose_exc: - await job.reset_budget_for_litellm_endusers() + await job.reset_budget_for_litellm_budget_table() await asyncio.sleep(0.1) # Verify exception logging assert mock_verbose_exc.call_count >= 1 @@ -949,3 +988,69 @@ async def test_service_logger_users_failure(): endusers_found_str = event_metadata.get("endusers_found", "") assert "user1" in endusers_found_str proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() + + +@pytest.mark.asyncio +async def test_reset_budget_for_litellm_team_members_called(): + """ + Test that when reset_budget_for_litellm_budget_table is called, + team members' budgets are also reset via reset_budget_for_litellm_team_members + """ + # Arrange + budget1 = LiteLLM_BudgetTableFull( + **{ + "budget_id": "budget1", + "max_budget": 100.0, + "budget_duration": "1d", + "created_at": datetime.now(timezone.utc) - timedelta(days=2), + } + ) + + enduser1 = {"user_id": "user1", "spend": 25.0, "budget_id": "budget1"} + + prisma_client = MagicMock() + + async def fake_get_data(*, table_name, query_type, **kwargs): + if table_name == "budget": + return [budget1] + elif table_name == "enduser": + return [enduser1] + return [] + + prisma_client.get_data = AsyncMock(side_effect=fake_get_data) + prisma_client.update_data = AsyncMock() + + # Mock the db.litellm_teammembership.update_many call + prisma_client.db = MagicMock() + prisma_client.db.litellm_teammembership = MagicMock() + prisma_client.db.litellm_teammembership.update_many = AsyncMock( + return_value={"count": 2} + ) + + proxy_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + + job = ResetBudgetJob(proxy_logging_obj, prisma_client) + + async def fake_reset_enduser(enduser): + enduser["spend"] = 0.0 + return enduser + + with patch.object( + ResetBudgetJob, + "_reset_budget_for_enduser", + side_effect=fake_reset_enduser, + ): + # Act + await job.reset_budget_for_litellm_budget_table() + + # Assert + # Verify that the team membership update was called + prisma_client.db.litellm_teammembership.update_many.assert_called_once() + + # Verify the call was made with correct parameters + call_args = prisma_client.db.litellm_teammembership.update_many.call_args + assert call_args.kwargs["where"]["budget_id"]["in"] == ["budget1"] + assert call_args.kwargs["data"]["spend"] == 0 diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 9763a15e58..a059a3adcb 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -3,6 +3,7 @@ import os import sys import time from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List import pytest @@ -16,19 +17,70 @@ from litellm.proxy.utils import ProxyLogging # Mock classes for testing +class MockLiteLLMTeamMembership: + async def update_many( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Dict[str, Any]: + # Mock the update_many method for litellm_teammembership + return {"count": 1} + + +class MockDB: + def __init__(self): + self.litellm_teammembership = MockLiteLLMTeamMembership() + + class MockPrismaClient: def __init__(self): - self.data = {"key": [], "user": [], "team": [], "budget": [], "enduser": []} - self.updated_data = { + self.data: Dict[str, List[Any]] = { "key": [], "user": [], "team": [], "budget": [], "enduser": [], } + self.updated_data: Dict[str, List[Any]] = { + "key": [], + "user": [], + "team": [], + "budget": [], + "enduser": [], + } + self.db = MockDB() async def get_data(self, table_name, query_type, **kwargs): - return self.data.get(table_name, []) + data = self.data.get(table_name, []) + + # Handle specific filtering for budget table queries + if table_name == "budget" and query_type == "find_all" and "reset_at" in kwargs: + # Return budgets that need to be reset (simulate expired budgets) + return [item for item in data if hasattr(item, "budget_reset_at")] + + # Handle specific filtering for enduser table queries + if ( + table_name == "enduser" + and query_type == "find_all" + and "budget_id_list" in kwargs + ): + budget_id_list = kwargs["budget_id_list"] + # Return endusers that match the budget IDs + return [ + item + for item in data + if hasattr(item, "litellm_budget_table") + and hasattr(item.litellm_budget_table, "budget_id") + and item.litellm_budget_table.budget_id in budget_id_list + ] + + # Handle key queries with expires and reset_at + if ( + table_name == "key" + and query_type == "find_all" + and ("expires" in kwargs or "reset_at" in kwargs) + ): + return [item for item in data if hasattr(item, "budget_reset_at")] + + return data async def update_data(self, query_type, data_list, table_name): self.updated_data[table_name] = data_list @@ -177,7 +229,7 @@ def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client): mock_prisma_client.data["enduser"] = [test_enduser] # Run the test - asyncio.run(reset_budget_job.reset_budget_for_litellm_endusers()) + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) # Verify results assert len(mock_prisma_client.updated_data["enduser"]) == 1 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 619739277a..c429662ec5 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -617,8 +617,8 @@ async def test_get_config_from_file(tmp_path, monkeypatch): async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): """ Test that _add_proxy_budget_to_db only creates a user and no keys are added. - - This validates that generate_key_helper_fn is called with table_name="user" + + This validates that generate_key_helper_fn is called with table_name="user" which should prevent key creation in LiteLLM_VerificationToken table. """ from unittest.mock import AsyncMock, patch @@ -629,31 +629,36 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): # Set up required litellm settings litellm.budget_duration = "30d" litellm.max_budget = 100.0 - + litellm_proxy_budget_name = "litellm-proxy-budget" # Mock generate_key_helper_fn to capture its call arguments - mock_generate_key_helper = AsyncMock(return_value={ - "user_id": litellm_proxy_budget_name, - "max_budget": 100.0, - "budget_duration": "30d", - "spend": 0, - "models": [], - }) - + mock_generate_key_helper = AsyncMock( + return_value={ + "user_id": litellm_proxy_budget_name, + "max_budget": 100.0, + "budget_duration": "30d", + "spend": 0, + "models": [], + } + ) + # Patch generate_key_helper_fn in proxy_server where it's being called from - with patch("litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper): + with patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper + ): # Call the function under test ProxyStartupEvent._add_proxy_budget_to_db(litellm_proxy_budget_name) - + # Allow async task to complete import asyncio + await asyncio.sleep(0.1) - + # Verify that generate_key_helper_fn was called mock_generate_key_helper.assert_called_once() call_args = mock_generate_key_helper.call_args - + # Verify critical parameters that prevent key creation assert call_args.kwargs["request_type"] == "user" assert call_args.kwargs["table_name"] == "user" @@ -682,38 +687,45 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): }, "model_list": [], "router_settings": {}, - "litellm_settings": {} + "litellm_settings": {}, } # Create temporary config file - with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: yaml.dump(test_config, f) config_file_path = f.name # Mock the get_instance_fn to return a mock handler mock_custom_handler = MagicMock() - + try: - with patch("litellm.proxy.proxy_server.get_instance_fn", return_value=mock_custom_handler) as mock_get_instance: + with patch( + "litellm.proxy.proxy_server.get_instance_fn", + return_value=mock_custom_handler, + ) as mock_get_instance: # Create ProxyConfig instance and load config proxy_config = ProxyConfig() # Create a mock router since load_config requires it mock_router = MagicMock() - await proxy_config.load_config(router=mock_router, config_file_path=config_file_path) - + await proxy_config.load_config( + router=mock_router, config_file_path=config_file_path + ) + # Verify get_instance_fn was called with correct parameters mock_get_instance.assert_called_with( value="custom_hooks.custom_ui_sso_hook.custom_ui_sso_sign_in_handler", - config_file_path=config_file_path + config_file_path=config_file_path, ) - + # Verify the global variable was set from litellm.proxy.proxy_server import user_custom_ui_sso_sign_in_handler + assert user_custom_ui_sso_sign_in_handler == mock_custom_handler - + finally: # Clean up temporary file import os + os.unlink(config_file_path) @@ -727,35 +739,41 @@ async def test_load_environment_variables_direct_and_os_environ(): from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() - + # Test config with both direct values and os.environ/ prefixed values test_config = { "environment_variables": { "DIRECT_VAR": "direct_value", "NUMERIC_VAR": 12345, "BOOL_VAR": True, - "SECRET_VAR": "os.environ/ACTUAL_SECRET_VAR" + "SECRET_VAR": "os.environ/ACTUAL_SECRET_VAR", } } - + # Mock get_secret_str to return a resolved value mock_secret_value = "resolved_secret_value" - - with patch("litellm.proxy.proxy_server.get_secret_str", return_value=mock_secret_value) as mock_get_secret: - with patch.dict(os.environ, {}, clear=False): # Don't clear existing env vars, just track changes + + with patch( + "litellm.proxy.proxy_server.get_secret_str", return_value=mock_secret_value + ) as mock_get_secret: + with patch.dict( + os.environ, {}, clear=False + ): # Don't clear existing env vars, just track changes # Call the method under test proxy_config._load_environment_variables(test_config) - + # Verify direct environment variables were set correctly assert os.environ["DIRECT_VAR"] == "direct_value" assert os.environ["NUMERIC_VAR"] == "12345" # Should be converted to string assert os.environ["BOOL_VAR"] == "True" # Should be converted to string - + # Verify os.environ/ prefixed variable was resolved and set assert os.environ["SECRET_VAR"] == mock_secret_value - + # Verify get_secret_str was called with the correct value - mock_get_secret.assert_called_once_with(secret_name="os.environ/ACTUAL_SECRET_VAR") + mock_get_secret.assert_called_once_with( + secret_name="os.environ/ACTUAL_SECRET_VAR" + ) @pytest.mark.asyncio @@ -768,56 +786,156 @@ async def test_load_environment_variables_litellm_license_and_edge_cases(): from litellm.proxy.proxy_server import ProxyConfig proxy_config = ProxyConfig() - + # Test Case 1: LITELLM_LICENSE in environment_variables test_config_with_license = { "environment_variables": { "LITELLM_LICENSE": "test_license_key", - "OTHER_VAR": "other_value" + "OTHER_VAR": "other_value", } } - + # Mock _license_check mock_license_check = MagicMock() mock_license_check.is_premium.return_value = True - + with patch("litellm.proxy.proxy_server._license_check", mock_license_check): with patch.dict(os.environ, {}, clear=False): # Call the method under test proxy_config._load_environment_variables(test_config_with_license) - + # Verify LITELLM_LICENSE was set in environment assert os.environ["LITELLM_LICENSE"] == "test_license_key" - + # Verify license check was updated assert mock_license_check.license_str == "test_license_key" mock_license_check.is_premium.assert_called_once() - + # Test Case 2: No environment_variables in config test_config_no_env_vars = {} - + # This should not raise any errors and should return without doing anything result = proxy_config._load_environment_variables(test_config_no_env_vars) assert result is None # Method returns None - + # Test Case 3: environment_variables is None test_config_none_env_vars = {"environment_variables": None} - + # This should not raise any errors and should return without doing anything result = proxy_config._load_environment_variables(test_config_none_env_vars) assert result is None # Method returns None - + # Test Case 4: os.environ/ prefix but get_secret_str returns None test_config_secret_none = { - "environment_variables": { - "FAILED_SECRET": "os.environ/NONEXISTENT_SECRET" - } + "environment_variables": {"FAILED_SECRET": "os.environ/NONEXISTENT_SECRET"} } - + with patch("litellm.proxy.proxy_server.get_secret_str", return_value=None): with patch.dict(os.environ, {}, clear=False): # Call the method under test proxy_config._load_environment_variables(test_config_secret_none) - + # Verify that the environment variable was not set when secret resolution fails assert "FAILED_SECRET" not in os.environ + + +@pytest.mark.asyncio +async def test_write_config_to_file(monkeypatch): + """ + Do not write config to file if store_model_in_db is True + """ + from unittest.mock import AsyncMock, MagicMock, mock_open, patch + + from litellm.proxy.proxy_server import ProxyConfig + + # Set store_model_in_db to True + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) + + # Mock prisma_client to not be None (so DB path is taken) + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Mock general_settings + mock_general_settings = {"store_model_in_db": True} + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", mock_general_settings + ) + + # Mock user_config_file_path + test_config_path = "/tmp/test_config.yaml" + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_config_file_path", test_config_path + ) + + proxy_config = ProxyConfig() + + # Mock the open function to track if file writing is attempted + mock_file_open = mock_open() + + with patch("builtins.open", mock_file_open), patch("yaml.dump") as mock_yaml_dump: + # Call save_config with test data + test_config = {"key": "value", "model_list": ["model1", "model2"]} + await proxy_config.save_config(new_config=test_config) + + # Verify that file was NOT opened for writing (since store_model_in_db=True) + mock_file_open.assert_not_called() + mock_yaml_dump.assert_not_called() + + # Verify that database insert was called instead + mock_prisma_client.insert_data.assert_called_once() + + # Verify the config passed to DB has model_list removed + call_args = mock_prisma_client.insert_data.call_args + assert call_args.kwargs["data"] == { + "key": "value" + } # model_list should be popped + assert call_args.kwargs["table_name"] == "config" + + +@pytest.mark.asyncio +async def test_write_config_to_file_when_store_model_in_db_false(monkeypatch): + """ + Test that config IS written to file when store_model_in_db is False + """ + from unittest.mock import AsyncMock, MagicMock, mock_open, patch + + from litellm.proxy.proxy_server import ProxyConfig + + # Set store_model_in_db to False + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + + # Mock prisma_client to be None (so file path is taken) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + + # Mock general_settings + mock_general_settings = {"store_model_in_db": False} + monkeypatch.setattr( + "litellm.proxy.proxy_server.general_settings", mock_general_settings + ) + + # Mock user_config_file_path + test_config_path = "/tmp/test_config.yaml" + monkeypatch.setattr( + "litellm.proxy.proxy_server.user_config_file_path", test_config_path + ) + + proxy_config = ProxyConfig() + + # Mock the open function and yaml.dump + mock_file_open = mock_open() + + with patch("builtins.open", mock_file_open), patch("yaml.dump") as mock_yaml_dump: + # Call save_config with test data + test_config = {"key": "value", "other_key": "other_value"} + await proxy_config.save_config(new_config=test_config) + + # Verify that file WAS opened for writing (since store_model_in_db=False) + mock_file_open.assert_called_once_with(f"{test_config_path}", "w") + + # Verify yaml.dump was called with the config + mock_yaml_dump.assert_called_once_with( + test_config, + mock_file_open.return_value.__enter__.return_value, + default_flow_style=False, + ) diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index db7045aa0d..98dbb589ef 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -11,7 +11,10 @@ sys.path.insert( from litellm.proxy._types import DefaultInternalUserParams, LitellmUserRoles from litellm.proxy.proxy_server import app -from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams, SSOConfig +from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + SSOConfig, +) client = TestClient(app) @@ -35,16 +38,14 @@ def mock_proxy_config(monkeypatch): "rpm_limit": 10, }, }, - "general_settings": { - "proxy_admin_email": "admin@example.com" - }, + "general_settings": {"proxy_admin_email": "admin@example.com"}, "environment_variables": { "GOOGLE_CLIENT_ID": "test_google_client_id", "GOOGLE_CLIENT_SECRET": "test_google_client_secret", "MICROSOFT_CLIENT_ID": "test_microsoft_client_id", "MICROSOFT_CLIENT_SECRET": "test_microsoft_client_secret", - "PROXY_BASE_URL": "https://example.com" - } + "PROXY_BASE_URL": "https://example.com", + }, } async def mock_get_config(): @@ -118,8 +119,10 @@ class TestProxySettingEndpoints: ): """Test updating the internal user settings""" # Mock litellm.default_internal_user_params + import litellm + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) monkeypatch.setattr(litellm, "default_internal_user_params", {}) # New settings to update @@ -132,6 +135,7 @@ class TestProxySettingEndpoints: response = client.patch("/update/internal_user_settings", json=new_settings) + print(response.text) assert response.status_code == 200 data = response.json() @@ -190,6 +194,7 @@ class TestProxySettingEndpoints: # Mock litellm.default_team_params import litellm + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True) monkeypatch.setattr(litellm, "default_team_params", {}) # New settings to update @@ -271,7 +276,7 @@ class TestProxySettingEndpoints: "microsoft_client_id": "new_microsoft_client_id", "microsoft_client_secret": "new_microsoft_client_secret", "proxy_base_url": "https://newexample.com", - "user_email": "newadmin@example.com" + "user_email": "newadmin@example.com", } response = client.patch("/update/sso_settings", json=new_sso_settings) @@ -286,17 +291,33 @@ class TestProxySettingEndpoints: # Verify settings were updated settings = data["settings"] assert settings["google_client_id"] == new_sso_settings["google_client_id"] - assert settings["google_client_secret"] == new_sso_settings["google_client_secret"] - assert settings["microsoft_client_id"] == new_sso_settings["microsoft_client_id"] - assert settings["microsoft_client_secret"] == new_sso_settings["microsoft_client_secret"] + assert ( + settings["google_client_secret"] == new_sso_settings["google_client_secret"] + ) + assert ( + settings["microsoft_client_id"] == new_sso_settings["microsoft_client_id"] + ) + assert ( + settings["microsoft_client_secret"] + == new_sso_settings["microsoft_client_secret"] + ) assert settings["proxy_base_url"] == new_sso_settings["proxy_base_url"] assert settings["user_email"] == new_sso_settings["user_email"] # Verify the config was updated updated_config = mock_proxy_config["config"] - assert updated_config["environment_variables"]["GOOGLE_CLIENT_ID"] == new_sso_settings["google_client_id"] - assert updated_config["environment_variables"]["GOOGLE_CLIENT_SECRET"] == new_sso_settings["google_client_secret"] - assert updated_config["general_settings"]["proxy_admin_email"] == new_sso_settings["user_email"] + assert ( + updated_config["environment_variables"]["GOOGLE_CLIENT_ID"] + == new_sso_settings["google_client_id"] + ) + assert ( + updated_config["environment_variables"]["GOOGLE_CLIENT_SECRET"] + == new_sso_settings["google_client_secret"] + ) + assert ( + updated_config["general_settings"]["proxy_admin_email"] + == new_sso_settings["user_email"] + ) # Verify save_config was called exactly once assert mock_proxy_config["save_call_count"]() == 1 diff --git a/ui/litellm-dashboard/src/components/SSOSettings.tsx b/ui/litellm-dashboard/src/components/SSOSettings.tsx index d414f6e54d..e21d7d1e00 100644 --- a/ui/litellm-dashboard/src/components/SSOSettings.tsx +++ b/ui/litellm-dashboard/src/components/SSOSettings.tsx @@ -81,7 +81,7 @@ const SSOSettings: React.FC = ({ accessToken, possibleUIRoles, setIsEditing(false); } catch (error) { console.error("Error updating SSO settings:", error); - message.error("Failed to update settings"); + message.error("Failed to update settings: " + error); } finally { setSaving(false); } diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 97940d19eb..ada75f9754 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -4836,7 +4836,7 @@ export const updateInternalUserSettings = async ( if (!response.ok) { const errorData = await response.text(); handleError(errorData); - throw new Error("Network response was not ok"); + throw new Error(errorData); } const data = await response.json();