diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index ca8c345f46..b2df9a1a2c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -24,13 +24,13 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, get_daily_activity_aggregated, ) -from litellm.proxy.auth.auth_checks import get_team_object, get_user_object from litellm.proxy.management_endpoints.common_utils import ( _is_user_team_admin, _user_has_admin_view, @@ -557,6 +557,18 @@ def get_team_from_list( return None +def _is_valid_user_id(user_id: str) -> bool: + """Validate that a decoded user_id is safe to use downstream.""" + MAX_USER_ID_LENGTH = 512 + if len(user_id) > MAX_USER_ID_LENGTH: + return False + # Reject null bytes and control characters (< 0x20) except space + for ch in user_id: + if ch == "\x00" or (ord(ch) < 0x20 and ch != " "): + return False + return True + + def get_user_id_from_request(request: Request) -> Optional[str]: """ Get the user id from the request @@ -573,7 +585,8 @@ def get_user_id_from_request(request: Request) -> Optional[str]: if match: # Use unquote instead of unquote_plus to preserve + characters raw_user_id = unquote(match.group(1)) - user_id = raw_user_id + if _is_valid_user_id(raw_user_id): + user_id = raw_user_id return user_id diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7d3d2ceb53..203344d452 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -480,11 +480,11 @@ from litellm.proxy.search_endpoints.search_tool_management import ( router as search_tool_management_router, ) from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router -from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.spend_tracking.spend_management_endpoints import ( router as spend_management_router, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.spend_tracking.vantage_endpoints import router as vantage_router from litellm.proxy.types_utils.utils import get_instance_fn from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( router as ui_crud_endpoints_router, @@ -2695,12 +2695,38 @@ class ProxyConfig: return search_tools_parsed if search_tools_parsed else None + # Environment variable keys that must not be overridden via config because + # they can alter process execution, library loading, or network routing. + _BLOCKED_ENV_KEYS: Set[str] = { + "PATH", + "LD_PRELOAD", + "LD_LIBRARY_PATH", + "DYLD_LIBRARY_PATH", + "DYLD_INSERT_LIBRARIES", + "PYTHONPATH", + "PYTHONSTARTUP", + "PYTHONHOME", + "HOME", + "USER", + "SHELL", + "LOGNAME", + "http_proxy", + "https_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + } + def _load_environment_variables(self, config: dict): ## ENVIRONMENT VARIABLES global premium_user environment_variables = config.get("environment_variables", None) if environment_variables: for key, value in environment_variables.items(): + if key in self._BLOCKED_ENV_KEYS: + verbose_proxy_logger.warning( + "Skipping blocked environment variable key: %s", key + ) + continue ######################################################### # handles this scenario: # ```yaml diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index e358cbe3be..0f90d236ae 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -85,6 +85,7 @@ async def test_ui_view_users_proxy_admin_no_org_filter(mocker): Proxy admin: find_many is called without organization_memberships in where. """ mock_prisma_client = mocker.MagicMock() + async def mock_find_many(*args, **kwargs): assert "organization_memberships" not in (kwargs.get("where") or {}) return [] @@ -327,6 +328,7 @@ async def test_ui_view_users_flag_on_team_admin_non_org_team_403(mocker): Flag ON, team admin for non-org team: returns 403. """ from fastapi import HTTPException + from litellm.proxy._types import LiteLLM_TeamTableCachedObj mock_prisma_client = mocker.MagicMock() @@ -372,9 +374,7 @@ async def test_ui_view_users_flag_on_team_admin_non_org_team_403(mocker): with pytest.raises(HTTPException) as exc_info: await ui_view_users( - user_api_key_dict=UserAPIKeyAuth( - user_id="team-admin-user", user_role=None - ), + user_api_key_dict=UserAPIKeyAuth(user_id="team-admin-user", user_role=None), user_id=None, user_email="u", team_id=tid, @@ -633,7 +633,9 @@ async def test_get_users_includes_timestamps(mocker): # Call get_users function directly with proxy admin auth admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) - response = await get_users(page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) + response = await get_users( + page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None + ) print("user /list response: ", response) @@ -855,7 +857,9 @@ async def test_new_user_non_admin_cannot_create_admin(mocker): # Verify the exception details assert exc_info.value.code == 403 or exc_info.value.code == "403" - assert "Only proxy admins can create administrative users" in str(exc_info.value.message) + assert "Only proxy admins can create administrative users" in str( + exc_info.value.message + ) assert "proxy_admin" in str(exc_info.value.message) assert "proxy_admin_viewer" in str(exc_info.value.message) assert str(LitellmUserRoles.PROXY_ADMIN) in str(exc_info.value.message) @@ -896,14 +900,14 @@ async def test_user_info_url_encoding_plus_character(mocker): # Mock the prisma client mock_prisma_client = mocker.MagicMock() - + # Create a real LiteLLM_UserTable instance (BaseModel) so isinstance check passes mock_user = LiteLLM_UserTable( user_id="machine-user+alp-air-admin-b58-b@tempus.com", user_email="machine-user+alp-air-admin-b58-b@tempus.com", teams=[], ) - + # Mock get_data to return user when called with user_id, empty list for keys async def mock_get_data(*args, **kwargs): if kwargs.get("table_name") == "key": @@ -913,7 +917,7 @@ async def test_user_info_url_encoding_plus_character(mocker): elif kwargs.get("user_id") is not None: return mock_user return None - + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) # Mock list_team to return None (patch it from where it's imported) @@ -941,7 +945,7 @@ async def test_user_info_url_encoding_plus_character(mocker): "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us ) expected_user_id = "machine-user+alp-air-admin-b58-b@tempus.com" - + response = await user_info( user_id=decoded_user_id, user_api_key_dict=mock_user_api_key_dict, @@ -955,7 +959,7 @@ async def test_user_info_url_encoding_plus_character(mocker): if call.kwargs.get("user_id") and not call.kwargs.get("table_name"): user_call = call break - + assert user_call is not None, "get_data should be called with user_id" assert user_call.kwargs["user_id"] == expected_user_id @@ -972,7 +976,7 @@ async def test_user_info_nonexistent_user(mocker): # Mock the prisma client mock_prisma_client = mocker.MagicMock() - + # Mock get_data to return None (user doesn't exist) async def mock_get_data(*args, **kwargs): if kwargs.get("table_name") == "key": @@ -980,7 +984,7 @@ async def test_user_info_nonexistent_user(mocker): elif kwargs.get("user_id") is not None: return None # User not found return None - + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) # Patch the prisma client import in the endpoint @@ -996,7 +1000,7 @@ async def test_user_info_nonexistent_user(mocker): # Call user_info function with a non-existent user_id nonexistent_user_id = "nonexistent-user@example.com" - + # Should raise ProxyException with 404 status code (HTTPException is converted by decorator) with pytest.raises(ProxyException) as exc_info: await user_info( @@ -1370,9 +1374,7 @@ async def test_check_duplicate_user_id(mocker): await _check_duplicate_user_id("existing-user-id", mock_prisma_client) assert exc_info.value.status_code == 409 - assert "User with id existing-user-id already exists" in str( - exc_info.value.detail - ) + assert "User with id existing-user-id already exists" in str(exc_info.value.detail) # No duplicate should pass async def mock_find_first_no_duplicate(*args, **kwargs): @@ -1393,7 +1395,7 @@ async def test_check_duplicate_user_id(mocker): def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): """ Test that _process_keys_for_user_info filters out keys with team_id='litellm-dashboard' - + UI session tokens (team_id='litellm-dashboard') should be excluded from user info responses to prevent confusion, as these are automatically created during dashboard login. """ @@ -1412,7 +1414,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): "user_id": "test-user", "key_alias": "dashboard-session-key", } - + mock_key_regular = MagicMock() mock_key_regular.model_dump.return_value = { "token": "sk-regular-token", @@ -1420,7 +1422,7 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): "user_id": "test-user", "key_alias": "regular-key", } - + mock_key_no_team = MagicMock() mock_key_no_team.model_dump.return_value = { "token": "sk-no-team-token", @@ -1446,20 +1448,24 @@ def test_process_keys_for_user_info_filters_dashboard_keys(monkeypatch): # Verify that dashboard key is filtered out assert len(result) == 2, "Should return 2 keys (dashboard key filtered out)" - + # Verify dashboard key is not in results result_team_ids = [key.get("team_id") for key in result] - assert UI_SESSION_TOKEN_TEAM_ID not in result_team_ids, "Dashboard key should be filtered out" - + assert ( + UI_SESSION_TOKEN_TEAM_ID not in result_team_ids + ), "Dashboard key should be filtered out" + # Verify regular keys are included assert "regular-team" in result_team_ids, "Regular team key should be included" assert None in result_team_ids, "No-team key should be included" - + # Verify the correct keys are returned result_tokens = [key.get("token") for key in result] assert "sk-regular-token" in result_tokens, "Regular key should be included" assert "sk-no-team-token" in result_tokens, "No-team key should be included" - assert "sk-dashboard-token" not in result_tokens, "Dashboard key should not be included" + assert ( + "sk-dashboard-token" not in result_tokens + ), "Dashboard key should not be included" def test_process_keys_for_user_info_handles_none_keys(monkeypatch): @@ -1558,7 +1564,13 @@ async def test_get_users_user_id_partial_match(mocker): admin_key = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) captured_where_conditions.clear() - await get_users(user_ids="test-user", page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) + await get_users( + user_ids="test-user", + page=1, + page_size=1, + user_api_key_dict=admin_key, + organization_ids=None, + ) assert "user_id" in captured_where_conditions assert "contains" in captured_where_conditions["user_id"] @@ -1566,7 +1578,13 @@ async def test_get_users_user_id_partial_match(mocker): assert captured_where_conditions["user_id"]["mode"] == "insensitive" captured_where_conditions.clear() - await get_users(user_ids="user1,user2,user3", page=1, page_size=1, user_api_key_dict=admin_key, organization_ids=None) + await get_users( + user_ids="user1,user2,user3", + page=1, + page_size=1, + user_api_key_dict=admin_key, + organization_ids=None, + ) assert "user_id" in captured_where_conditions assert "in" in captured_where_conditions["user_id"] @@ -1578,7 +1596,7 @@ def test_update_internal_user_params_reset_max_budget_with_none(): Test that _update_internal_user_params allows setting max_budget to None. This verifies the fix for unsetting/resetting the budget to unlimited. """ - + # Case 1: max_budget is explicitly None in the input dictionary data_json = {"max_budget": None, "user_id": "test_user"} data = UpdateUserRequest(max_budget=None, user_id="test_user") @@ -1610,7 +1628,7 @@ def test_update_internal_user_params_ignores_other_nones(): def test_update_internal_user_params_keeps_original_max_budget_when_not_provided(): """ - Test that _update_internal_user_params does not include max_budget + Test that _update_internal_user_params does not include max_budget when it's not provided in the request (should keep original value). """ # Create test data without max_budget @@ -1631,7 +1649,7 @@ def test_generate_request_base_validator(): Test that GenerateRequestBase validator converts empty string to None for max_budget """ from litellm.proxy._types import GenerateRequestBase - + # Test with empty string req = GenerateRequestBase(max_budget="") assert req.max_budget is None @@ -1662,9 +1680,7 @@ async def test_get_user_daily_activity_non_admin_cannot_view_other_users(monkeyp # Mock the prisma client so the DB-not-connected check passes mock_prisma_client = MagicMock() - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", mock_prisma_client - ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Non-admin caller non_admin_key_dict = UserAPIKeyAuth( @@ -1731,9 +1747,7 @@ async def test_get_user_daily_activity_aggregated_admin_global_view(monkeypatch) # Mock the prisma client mock_prisma_client = MagicMock() - monkeypatch.setattr( - "litellm.proxy.proxy_server.prisma_client", mock_prisma_client - ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) # Mock the downstream helper so we don't need a real DB mock_response = MagicMock() @@ -1846,7 +1860,9 @@ async def test_delete_user_cleans_up_created_by_invitation_links(mocker): call_kwargs = mock_prisma_client.db.litellm_invitationlink.delete_many.call_args where_clause = call_kwargs.kwargs.get("where") or call_kwargs[1].get("where") - assert "OR" in where_clause, "Should use OR to match user_id, created_by, and updated_by" + assert ( + "OR" in where_clause + ), "Should use OR to match user_id, created_by, and updated_by" or_conditions = where_clause["OR"] assert len(or_conditions) == 3, "Should have 3 OR conditions" @@ -2188,9 +2204,20 @@ async def test_user_info_v2_response_shape(mocker): # Verify all expected fields are present response_dict = response.model_dump() expected_fields = { - "user_id", "user_email", "user_alias", "user_role", "spend", - "max_budget", "models", "budget_duration", "budget_reset_at", - "metadata", "created_at", "updated_at", "sso_user_id", "teams", + "user_id", + "user_email", + "user_alias", + "user_role", + "spend", + "max_budget", + "models", + "budget_duration", + "budget_reset_at", + "metadata", + "created_at", + "updated_at", + "sso_user_id", + "teams", } assert set(response_dict.keys()) == expected_fields @@ -2418,4 +2445,74 @@ async def test_user_info_v2_url_encoding_plus_character(mocker): ) assert isinstance(response, UserInfoV2Response) - assert response.user_id == expected_user_id \ No newline at end of file + assert response.user_id == expected_user_id + + +class TestGetUserIdFromRequestValidation: + """Tests for user_id input validation in get_user_id_from_request.""" + + def _make_request(self, query_string: str): + from unittest.mock import MagicMock + + from starlette.requests import Request + + request = MagicMock(spec=Request) + request.url.query = query_string + return request + + def test_valid_uuid(self): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_id_from_request, + ) + + request = self._make_request("user_id=550e8400-e29b-41d4-a716-446655440000") + result = get_user_id_from_request(request) + assert result == "550e8400-e29b-41d4-a716-446655440000" + + def test_valid_email(self): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_id_from_request, + ) + + request = self._make_request("user_id=user%40example.com") + result = get_user_id_from_request(request) + assert result == "user@example.com" + + def test_rejects_overlong_user_id(self): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_id_from_request, + ) + + long_id = "a" * 513 + request = self._make_request(f"user_id={long_id}") + result = get_user_id_from_request(request) + assert result is None + + def test_rejects_null_byte(self): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_id_from_request, + ) + + request = self._make_request("user_id=admin%00evil") + result = get_user_id_from_request(request) + assert result is None + + def test_rejects_control_characters(self): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_id_from_request, + ) + + # Tab character (0x09) + request = self._make_request("user_id=admin%09evil") + result = get_user_id_from_request(request) + assert result is None + + def test_allows_512_char_user_id(self): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + get_user_id_from_request, + ) + + exact_id = "a" * 512 + request = self._make_request(f"user_id={exact_id}") + result = get_user_id_from_request(request) + assert result == exact_id diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index bd6162f225..334cd7395d 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -104,10 +104,7 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) assert response.status_code == 200 - assert ( - response.json() - == {"redirect_url": "http://testserver/ui/?login=success"} - ) + assert response.json() == {"redirect_url": "http://testserver/ui/?login=success"} assert response.cookies.get("token") == "signed-token" mock_authenticate_user.assert_awaited_once_with( @@ -516,15 +513,11 @@ def test_restructure_ui_html_files_handles_nested_routes(tmp_path): assert (ui_root / "home" / "index.html").read_text() == "home" assert not (ui_root / "mcp" / "oauth" / "callback.html").exists() assert ( - (ui_root / "mcp" / "oauth" / "callback" / "index.html").read_text() - == "callback" - ) + ui_root / "mcp" / "oauth" / "callback" / "index.html" + ).read_text() == "callback" assert (ui_root / "existing" / "index.html").read_text() == "keep" assert (ui_root / "_next" / "ignore.html").read_text() == "asset" - assert ( - (ui_root / "litellm-asset-prefix" / "ignore.html").read_text() - == "asset" - ) + assert (ui_root / "litellm-asset-prefix" / "ignore.html").read_text() == "asset" def test_ui_extensionless_route_requires_restructure(tmp_path): @@ -541,9 +534,7 @@ def test_ui_extensionless_route_requires_restructure(tmp_path): (ui_root / "login.html").write_text("login") fastapi_app = FastAPI() - fastapi_app.mount( - "/ui", StaticFiles(directory=str(ui_root), html=True), name="ui" - ) + fastapi_app.mount("/ui", StaticFiles(directory=str(ui_root), html=True), name="ui") client = TestClient(fastapi_app) assert client.get("/ui/login.html").status_code == 200 @@ -564,37 +555,37 @@ def test_restructure_always_happens(monkeypatch): """ # Test Case 1: is_non_root is True - restructuring happens in /var/lib/litellm/ui monkeypatch.setenv("LITELLM_NON_ROOT", "true") - + runtime_ui_path = "/var/lib/litellm/ui" packaged_ui_path = "/some/packaged/ui/path" - + # Simulate the logic from proxy_server.py is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true" if is_non_root: ui_path = runtime_ui_path else: ui_path = packaged_ui_path - + # Restructuring always happens now, regardless of ui_path vs packaged_ui_path should_restructure = True - + assert is_non_root is True assert should_restructure is True assert ui_path == runtime_ui_path - + # Test Case 2: is_non_root is False - restructuring happens directly in packaged_ui_path monkeypatch.delenv("LITELLM_NON_ROOT", raising=False) - + # Simulate the logic from proxy_server.py is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true" if is_non_root: ui_path = runtime_ui_path else: ui_path = packaged_ui_path - + # Restructuring always happens now, even when ui_path == packaged_ui_path should_restructure = True - + assert is_non_root is False assert should_restructure is True assert ui_path == packaged_ui_path @@ -691,9 +682,7 @@ def test_update_config_fields_deep_merge_db_wins(): "hidden": True, }, # Demonstrate that None values from DB are skipped (preserve existing) - "legacy-sonnet": { - "hidden": None # should not clobber current True - }, + "legacy-sonnet": {"hidden": None}, # should not clobber current True } } @@ -743,9 +732,7 @@ def test_get_config_custom_callback_api_env_vars(monkeypatch): mock_router = MagicMock() mock_router.get_settings.return_value = {} monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) - monkeypatch.setattr( - proxy_config, "get_config", AsyncMock(return_value=config_data) - ) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) # Bypass auth dependency original_overrides = app.dependency_overrides.copy() @@ -923,7 +910,9 @@ def test_embedding_input_array_of_tokens(client_no_auth): assert response.status_code == 200 result = response.json() print(len(result["data"][0]["embedding"])) - assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so + assert ( + len(result["data"][0]["embedding"]) > 10 + ) # this usually has len==1536 so except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") @@ -1180,7 +1169,6 @@ async def test_delete_deployment_type_mismatch(): with patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), patch( "litellm.proxy.proxy_server.user_config_file_path", "test_config.yaml" ): - # Call the function under test deleted_count = await pc._delete_deployment(db_models=[]) @@ -1322,6 +1310,7 @@ def test_normalize_datetime_for_sorting(): # Test Case 6: Timezone-aware datetime object (non-UTC) from datetime import timedelta + aware_dt = datetime(2024, 1, 15, 10, 30, 0, tzinfo=timezone(timedelta(hours=5))) result = _normalize_datetime_for_sorting(aware_dt) assert result is not None @@ -1574,6 +1563,70 @@ async def test_load_environment_variables_litellm_license_and_edge_cases(): assert "FAILED_SECRET" not in os.environ +@pytest.mark.asyncio +async def test_load_environment_variables_blocks_dangerous_keys(): + """ + Test that _load_environment_variables rejects dangerous env var keys + like PATH, LD_PRELOAD, PYTHONPATH, etc. + """ + import logging + + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + original_path = os.environ.get("PATH", "") + + test_config = { + "environment_variables": { + "PATH": "/tmp/evil", + "LD_PRELOAD": "/tmp/evil.so", + "PYTHONPATH": "/tmp/evil", + "SAFE_CUSTOM_VAR": "safe_value", + } + } + + with patch.dict(os.environ, {}, clear=False): + proxy_config._load_environment_variables(test_config) + + # Blocked keys should not be set to the attacker value + assert os.environ.get("PATH") != "/tmp/evil" + assert ( + "LD_PRELOAD" not in os.environ or os.environ["LD_PRELOAD"] != "/tmp/evil.so" + ) + assert os.environ.get("PYTHONPATH") != "/tmp/evil" + + # Safe keys should still be set + assert os.environ["SAFE_CUSTOM_VAR"] == "safe_value" + + +@pytest.mark.asyncio +async def test_load_environment_variables_blocks_proxy_keys(): + """ + Test that _load_environment_variables rejects proxy-related env var keys. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + test_config = { + "environment_variables": { + "HTTP_PROXY": "http://evil-proxy:8080", + "HTTPS_PROXY": "http://evil-proxy:8080", + "http_proxy": "http://evil-proxy:8080", + "https_proxy": "http://evil-proxy:8080", + } + } + + with patch.dict(os.environ, {}, clear=False): + proxy_config._load_environment_variables(test_config) + + assert os.environ.get("HTTP_PROXY") != "http://evil-proxy:8080" + assert os.environ.get("HTTPS_PROXY") != "http://evil-proxy:8080" + assert os.environ.get("http_proxy") != "http://evil-proxy:8080" + assert os.environ.get("https_proxy") != "http://evil-proxy:8080" + + @pytest.mark.asyncio async def test_write_config_to_file(monkeypatch): """ @@ -1882,7 +1935,6 @@ async def test_chat_completion_result_no_nested_none_values(): "litellm.proxy.proxy_server.ProxyBaseLLMRequestProcessing", return_value=mock_base_processor, ): - # Call the chat_completion function result = await chat_completion( request=mock_request, @@ -2027,9 +2079,7 @@ class TestPriceDataReloadAPI: # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: - mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=None - ) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client_with_auth.post("/reload/model_cost_map") @@ -2372,8 +2422,13 @@ class TestPriceDataReloadIntegration: with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # Simulate existing config with a schedule mock_existing = MagicMock() - mock_existing.param_value = {"interval_hours": 12, "force_reload": False} - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_existing) + mock_existing.param_value = { + "interval_hours": 12, + "force_reload": False, + } + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=mock_existing + ) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client.post("/reload/model_cost_map") @@ -2415,7 +2470,9 @@ class TestPriceDataReloadIntegration: ) as mock_reload: mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} - asyncio.run(proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma)) + asyncio.run( + proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma) + ) # Verify the upsert update branch preserves interval_hours mock_prisma.db.litellm_config.upsert.assert_called() @@ -2456,7 +2513,9 @@ class TestPriceDataReloadIntegration: # Simulate existing config with a schedule mock_existing = MagicMock() mock_existing.param_value = {"interval_hours": 8, "force_reload": False} - mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_existing) + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=mock_existing + ) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=None) response = client.post("/reload/anthropic_beta_headers") @@ -2821,7 +2880,7 @@ async def test_model_info_v1_oci_secrets_not_leaked(): mock_user_api_key_dict.api_key = "test-key" mock_user_api_key_dict.team_models = [] mock_user_api_key_dict.models = ["oci-grok-test"] - + # Mock model data with OCI sensitive information mock_model_data = { "model_name": "oci-grok-test", @@ -2834,59 +2893,73 @@ async def test_model_info_v1_oci_secrets_not_leaked(): "oci_tenancy": "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk", "oci_key_file": "/path/to/oci_api_key.pem", "oci_compartment_id": "ocid1.compartment.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk", - "drop_params": True + "drop_params": True, }, - "model_info": { - "mode": "completion", - "id": "test-model-id" - } + "model_info": {"mode": "completion", "id": "test-model-id"}, } - + # Mock the llm_router to return our test data mock_router = MagicMock() mock_router.get_model_names.return_value = ["oci-grok-test"] mock_router.get_model_access_groups.return_value = {} mock_router.get_model_list.return_value = [mock_model_data] - + # Mock global variables - with patch("litellm.proxy.proxy_server.llm_router", mock_router), \ - patch("litellm.proxy.proxy_server.llm_model_list", [mock_model_data]), \ - patch("litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False}), \ - patch("litellm.proxy.proxy_server.user_model", None): - + with patch("litellm.proxy.proxy_server.llm_router", mock_router), patch( + "litellm.proxy.proxy_server.llm_model_list", [mock_model_data] + ), patch( + "litellm.proxy.proxy_server.general_settings", {"infer_model_from_keys": False} + ), patch( + "litellm.proxy.proxy_server.user_model", None + ): # Call the model_info_v1 endpoint result = await model_info_v1( - user_api_key_dict=mock_user_api_key_dict, - litellm_model_id=None + user_api_key_dict=mock_user_api_key_dict, litellm_model_id=None ) - + # Verify the result structure assert "data" in result assert len(result["data"]) == 1 - + model_info = result["data"][0] litellm_params = model_info["litellm_params"] - + # Verify that sensitive OCI fields are masked assert "****" in litellm_params["oci_key"], "oci_key should be masked" - assert "****" in litellm_params["oci_fingerprint"], "oci_fingerprint should be masked" + assert ( + "****" in litellm_params["oci_fingerprint"] + ), "oci_fingerprint should be masked" assert "****" in litellm_params["oci_tenancy"], "oci_tenancy should be masked" assert "****" in litellm_params["oci_key_file"], "oci_key_file should be masked" - + # Verify that non-sensitive fields are NOT masked - assert litellm_params["model"] == "oci/xai.grok-4", "model field should not be masked" - assert litellm_params["oci_region"] == "us-phoenix-1", "oci_region should not be masked" + assert ( + litellm_params["model"] == "oci/xai.grok-4" + ), "model field should not be masked" + assert ( + litellm_params["oci_region"] == "us-phoenix-1" + ), "oci_region should not be masked" assert litellm_params["drop_params"] is True, "drop_params should not be masked" - + # Verify the model field specifically is not masked (this was the original issue) - assert "****" not in litellm_params["model"], "model field should never be masked" - assert litellm_params["model"].startswith("oci/"), "model should retain its full value" - + assert ( + "****" not in litellm_params["model"] + ), "model field should never be masked" + assert litellm_params["model"].startswith( + "oci/" + ), "model should retain its full value" + # Verify that actual secret values are not present in the response result_str = str(result) - assert "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str + assert ( + "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" + not in result_str + ) assert "aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00" not in result_str - assert "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str + assert ( + "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" + not in result_str + ) assert "/path/to/oci_api_key.pem" not in result_str @@ -2898,17 +2971,17 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): from unittest.mock import MagicMock, patch from litellm.proxy.proxy_server import ProxyConfig - + proxy_config = ProxyConfig() - + # Mock the callback manager mock_callback_manager = MagicMock() - + with patch("litellm.proxy.proxy_server.litellm") as mock_litellm: # Set up mock litellm attributes mock_litellm._known_custom_logger_compatible_callbacks = [] mock_litellm.logging_callback_manager = mock_callback_manager - + # Test Case 1: Add success callback mock_success_callbacks = [] proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( @@ -2916,9 +2989,11 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): event_types=["success"], existing_callbacks=mock_success_callbacks, ) - mock_callback_manager.add_litellm_success_callback.assert_called_once_with("prometheus") + mock_callback_manager.add_litellm_success_callback.assert_called_once_with( + "prometheus" + ) mock_callback_manager.reset_mock() - + # Test Case 2: Add failure callback mock_failure_callbacks = [] proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( @@ -2926,9 +3001,11 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): event_types=["failure"], existing_callbacks=mock_failure_callbacks, ) - mock_callback_manager.add_litellm_failure_callback.assert_called_once_with("langfuse") + mock_callback_manager.add_litellm_failure_callback.assert_called_once_with( + "langfuse" + ) mock_callback_manager.reset_mock() - + # Test Case 3: Add callback for both success and failure mock_callbacks = [] proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( @@ -2938,7 +3015,7 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): ) mock_callback_manager.add_litellm_callback.assert_called_once_with("s3") mock_callback_manager.reset_mock() - + # Test Case 4: Don't add callback if it already exists existing_callbacks_with_item = ["prometheus"] proxy_config._add_callback_from_db_to_in_memory_litellm_callbacks( @@ -2952,7 +3029,7 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): def test_should_load_db_object_with_supported_db_objects(): """ Test _should_load_db_object method with supported_db_objects configuration. - + Verifies that when supported_db_objects is set, only specified object types are loaded from the database. """ @@ -3056,8 +3133,12 @@ async def test_tag_cache_update_called(): "spend": 10.0, } - with patch.object(cache, "async_get_cache", new=AsyncMock(return_value=mock_tag_obj)) as mock_get_cache: - with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: + with patch.object( + cache, "async_get_cache", new=AsyncMock(return_value=mock_tag_obj) + ) as mock_get_cache: + with patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock() + ) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -3108,8 +3189,12 @@ async def test_tag_cache_update_multiple_tags(): return mock_tag2_obj return None - with patch.object(cache, "async_get_cache", new=AsyncMock(side_effect=mock_get_cache_side_effect)) as mock_get_cache: - with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: + with patch.object( + cache, "async_get_cache", new=AsyncMock(side_effect=mock_get_cache_side_effect) + ) as mock_get_cache: + with patch.object( + cache, "async_set_cache_pipeline", new=AsyncMock() + ) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -3130,7 +3215,9 @@ async def test_tag_cache_update_multiple_tags(): assert len(cache_list) == 2 - tag_updates = {cache_key: cache_value for cache_key, cache_value in cache_list} + tag_updates = { + cache_key: cache_value for cache_key, cache_value in cache_list + } assert "tag:tag1" in tag_updates assert "tag:tag2" in tag_updates assert tag_updates["tag:tag1"]["spend"] == 15.0 @@ -3250,7 +3337,9 @@ async def test_init_sso_settings_in_db_error_handling(): assert True except Exception as e: # The exception should be caught and logged, not propagated - pytest.fail(f"Exception should have been caught and logged, but was raised: {e}") + pytest.fail( + f"Exception should have been caught and logged, but was raised: {e}" + ) @pytest.mark.asyncio @@ -3353,11 +3442,15 @@ def test_get_prompt_spec_for_db_prompt_with_versions(): } # Test version 1 - prompt_spec_v1 = proxy_config._get_prompt_spec_for_db_prompt(db_prompt=mock_prompt_v1) + prompt_spec_v1 = proxy_config._get_prompt_spec_for_db_prompt( + db_prompt=mock_prompt_v1 + ) assert prompt_spec_v1.prompt_id == "chat_prompt.v1" # Test version 2 - prompt_spec_v2 = proxy_config._get_prompt_spec_for_db_prompt(db_prompt=mock_prompt_v2) + prompt_spec_v2 = proxy_config._get_prompt_spec_for_db_prompt( + db_prompt=mock_prompt_v2 + ) assert prompt_spec_v2.prompt_id == "chat_prompt.v2" @@ -3372,15 +3465,15 @@ def test_root_redirect_when_docs_url_not_root_and_redirect_url_set(monkeypatch): config_fp = f"{filepath}/test_configs/test_config_no_auth.yaml" # Ensure docs are mounted on a non-root path to trigger redirect logic monkeypatch.setenv("DOCS_URL", "/docs") - + test_redirect_url = "/ui" monkeypatch.setenv("ROOT_REDIRECT_URL", test_redirect_url) - + asyncio.run(initialize(config=config_fp, debug=True)) - + docs_url = _get_docs_url() root_redirect_url = os.getenv("ROOT_REDIRECT_URL") - + # Remove any existing "/" route that might interfere routes_to_remove = [] for route in app.routes: @@ -3389,16 +3482,17 @@ def test_root_redirect_when_docs_url_not_root_and_redirect_url_set(monkeypatch): routes_to_remove.append(route) elif not hasattr(route, "methods"): # Catch-all routes routes_to_remove.append(route) - + for route in routes_to_remove: app.routes.remove(route) - + # Add the redirect route if conditions are met (matching the actual implementation) if docs_url != "/" and root_redirect_url: + @app.get("/", include_in_schema=False) async def root_redirect(): return RedirectResponse(url=root_redirect_url) - + client = TestClient(app) response = client.get("/", follow_redirects=False) assert response.status_code == 307 @@ -3422,12 +3516,13 @@ async def test_get_image_non_root_uses_var_lib_assets_dir(monkeypatch): def exists_side_effect(path): return False if path == "/var/lib/litellm/assets" else True - with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ - patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ - patch("litellm.proxy.proxy_server.os.access", return_value=True), \ - patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ - patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: - + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, patch( + "litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect + ), patch("litellm.proxy.proxy_server.os.access", return_value=True), patch( + "litellm.proxy.proxy_server.os.getenv" + ) as mock_getenv, patch( + "litellm.proxy.proxy_server.FileResponse" + ) as mock_file_response: # Setup mock_getenv to return empty string for UI_LOGO_PATH def getenv_side_effect(key, default=""): if key == "UI_LOGO_PATH": @@ -3471,12 +3566,13 @@ async def test_get_image_non_root_fallback_to_default_logo(monkeypatch): return True # Mock os.path operations - with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ - patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ - patch("litellm.proxy.proxy_server.os.access", return_value=True), \ - patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ - patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: - + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, patch( + "litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect + ), patch("litellm.proxy.proxy_server.os.access", return_value=True), patch( + "litellm.proxy.proxy_server.os.getenv" + ) as mock_getenv, patch( + "litellm.proxy.proxy_server.FileResponse" + ) as mock_file_response: # Setup mock_getenv def getenv_side_effect(key, default=""): if key == "UI_LOGO_PATH": @@ -3495,8 +3591,9 @@ async def test_get_image_non_root_fallback_to_default_logo(monkeypatch): # Verify that exists was called to check /var/lib/litellm/assets/logo.jpg assets_logo_path = "/var/lib/litellm/assets/logo.jpg" - assert any(assets_logo_path in str(call) for call in exists_calls), \ - f"Should check if {assets_logo_path} exists" + assert any( + assets_logo_path in str(call) for call in exists_calls + ), f"Should check if {assets_logo_path} exists" # Verify FileResponse was called (with fallback logo) assert mock_file_response.called, "FileResponse should be called" @@ -3516,11 +3613,11 @@ async def test_get_image_root_case_uses_current_dir(monkeypatch): monkeypatch.delenv("UI_LOGO_PATH", raising=False) # Mock os.path operations - with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ - patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), \ - patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ - patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: - + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, patch( + "litellm.proxy.proxy_server.os.path.exists", return_value=True + ), patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, patch( + "litellm.proxy.proxy_server.FileResponse" + ) as mock_file_response: # Setup mock_getenv def getenv_side_effect(key, default=""): if key == "UI_LOGO_PATH": @@ -3536,10 +3633,13 @@ async def test_get_image_root_case_uses_current_dir(monkeypatch): # Verify makedirs was NOT called with /var/lib/litellm/assets (should not create it for root case) var_lib_assets_calls = [ - call for call in mock_makedirs.call_args_list + call + for call in mock_makedirs.call_args_list if "/var/lib/litellm/assets" in str(call) ] - assert len(var_lib_assets_calls) == 0, "Should not create /var/lib/litellm/assets for root case" + assert ( + len(var_lib_assets_calls) == 0 + ), "Should not create /var/lib/litellm/assets for root case" # Verify FileResponse was called assert mock_file_response.called, "FileResponse should be called" @@ -3569,13 +3669,14 @@ async def test_get_image_custom_local_logo_bypasses_cache(monkeypatch): calls_to_file_response.append(path) return MagicMock() - with patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), \ - patch("litellm.proxy.proxy_server.os.access", return_value=True), \ - patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response): - + with patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), patch( + "litellm.proxy.proxy_server.os.access", return_value=True + ), patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response): await get_image() - assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" + assert ( + len(calls_to_file_response) == 1 + ), "FileResponse should be called exactly once" assert calls_to_file_response[0] == "/app/custom_logo.jpg", ( f"Expected custom logo path, got {calls_to_file_response[0]}. " "A stale cached_logo.jpg may have been returned instead." @@ -3602,17 +3703,18 @@ async def test_get_image_default_logo_still_uses_cache(monkeypatch): calls_to_file_response.append(path) return MagicMock() - with patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), \ - patch("litellm.proxy.proxy_server.os.access", return_value=True), \ - patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response): - + with patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), patch( + "litellm.proxy.proxy_server.os.access", return_value=True + ), patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response): await get_image() - assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" + assert ( + len(calls_to_file_response) == 1 + ), "FileResponse should be called exactly once" served_path = calls_to_file_response[0] - assert served_path.endswith("cached_logo.jpg"), ( - f"Expected cached_logo.jpg for default logo, got {served_path}" - ) + assert served_path.endswith( + "cached_logo.jpg" + ), f"Expected cached_logo.jpg for default logo, got {served_path}" @pytest.mark.asyncio @@ -3641,20 +3743,23 @@ async def test_get_image_custom_logo_missing_falls_through_to_default(monkeypatc return False return True - with patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ - patch("litellm.proxy.proxy_server.os.access", return_value=True), \ - patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response): - + with patch( + "litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect + ), patch("litellm.proxy.proxy_server.os.access", return_value=True), patch( + "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response + ): await get_image() - assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" + assert ( + len(calls_to_file_response) == 1 + ), "FileResponse should be called exactly once" served_path = calls_to_file_response[0] - assert served_path != "/app/nonexistent_logo.jpg", ( - "Should not attempt to serve a non-existent custom logo" - ) - assert served_path.endswith("cached_logo.jpg"), ( - f"Expected fallback to cached_logo.jpg, got {served_path}" - ) + assert ( + served_path != "/app/nonexistent_logo.jpg" + ), "Should not attempt to serve a non-existent custom logo" + assert served_path.endswith( + "cached_logo.jpg" + ), f"Expected fallback to cached_logo.jpg, got {served_path}" @pytest.mark.asyncio @@ -3686,20 +3791,23 @@ async def test_get_image_custom_logo_missing_no_cache_serves_default(monkeypatch return False return True - with patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ - patch("litellm.proxy.proxy_server.os.access", return_value=True), \ - patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response): - + with patch( + "litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect + ), patch("litellm.proxy.proxy_server.os.access", return_value=True), patch( + "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response + ): await get_image() - assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" + assert ( + len(calls_to_file_response) == 1 + ), "FileResponse should be called exactly once" served_path = calls_to_file_response[0] - assert served_path != "/app/nonexistent_logo.jpg", ( - "Should not attempt to serve a non-existent custom logo" - ) - assert served_path.endswith("logo.jpg"), ( - f"Expected fallback to default logo.jpg, got {served_path}" - ) + assert ( + served_path != "/app/nonexistent_logo.jpg" + ), "Should not attempt to serve a non-existent custom logo" + assert served_path.endswith( + "logo.jpg" + ), f"Expected fallback to default logo.jpg, got {served_path}" def test_get_config_normalizes_string_callbacks(monkeypatch): @@ -3721,9 +3829,7 @@ def test_get_config_normalizes_string_callbacks(monkeypatch): mock_router = MagicMock() mock_router.get_settings.return_value = {} monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) - monkeypatch.setattr( - proxy_config, "get_config", AsyncMock(return_value=config_data) - ) + monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) original_overrides = app.dependency_overrides.copy() app.dependency_overrides[user_api_key_auth] = lambda: MagicMock() @@ -4127,9 +4233,9 @@ async def test_update_general_settings_store_model_in_db_false(): proxy_config = ProxyConfig() - with patch( - "litellm.proxy.proxy_server.store_model_in_db", True - ), patch("litellm.proxy.proxy_server.general_settings", {}): + with patch("litellm.proxy.proxy_server.store_model_in_db", True), patch( + "litellm.proxy.proxy_server.general_settings", {} + ): await proxy_config._update_general_settings( db_general_settings={"store_model_in_db": False} ) @@ -4150,9 +4256,9 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): proxy_config = ProxyConfig() # Test "true" string - with patch( - "litellm.proxy.proxy_server.store_model_in_db", False - ), patch("litellm.proxy.proxy_server.general_settings", {}): + with patch("litellm.proxy.proxy_server.store_model_in_db", False), patch( + "litellm.proxy.proxy_server.general_settings", {} + ): await proxy_config._update_general_settings( db_general_settings={"store_model_in_db": "true"} ) @@ -4161,9 +4267,9 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): assert ps.store_model_in_db is True # Test "True" string - with patch( - "litellm.proxy.proxy_server.store_model_in_db", False - ), patch("litellm.proxy.proxy_server.general_settings", {}): + with patch("litellm.proxy.proxy_server.store_model_in_db", False), patch( + "litellm.proxy.proxy_server.general_settings", {} + ): await proxy_config._update_general_settings( db_general_settings={"store_model_in_db": "True"} ) @@ -4172,9 +4278,9 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): assert ps.store_model_in_db is True # Test "false" string - with patch( - "litellm.proxy.proxy_server.store_model_in_db", True - ), patch("litellm.proxy.proxy_server.general_settings", {}): + with patch("litellm.proxy.proxy_server.store_model_in_db", True), patch( + "litellm.proxy.proxy_server.general_settings", {} + ): await proxy_config._update_general_settings( db_general_settings={"store_model_in_db": "false"} ) @@ -4194,9 +4300,9 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): proxy_config = ProxyConfig() # When current is True and DB sends None, should stay True - with patch( - "litellm.proxy.proxy_server.store_model_in_db", True - ), patch("litellm.proxy.proxy_server.general_settings", {}): + with patch("litellm.proxy.proxy_server.store_model_in_db", True), patch( + "litellm.proxy.proxy_server.general_settings", {} + ): await proxy_config._update_general_settings( db_general_settings={"store_model_in_db": None} ) @@ -4205,9 +4311,9 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): assert ps.store_model_in_db is True # When current is False and DB sends None, should stay False - with patch( - "litellm.proxy.proxy_server.store_model_in_db", False - ), patch("litellm.proxy.proxy_server.general_settings", {}): + with patch("litellm.proxy.proxy_server.store_model_in_db", False), patch( + "litellm.proxy.proxy_server.general_settings", {} + ): await proxy_config._update_general_settings( db_general_settings={"store_model_in_db": None} ) @@ -4238,13 +4344,9 @@ async def test_store_model_in_db_db_override_when_config_false(): mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_config = AsyncMock() - with patch( - "litellm.proxy.proxy_server.proxy_config", mock_proxy_config - ), patch( + with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch( "litellm.proxy.proxy_server.store_model_in_db", False - ), patch( - "litellm.proxy.proxy_server.get_secret_bool", return_value=False - ): + ), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=False): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, prisma_client=mock_prisma_client, @@ -4282,13 +4384,9 @@ async def test_store_model_in_db_db_check_skipped_when_already_true(monkeypatch) mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_config = AsyncMock() - with patch( - "litellm.proxy.proxy_server.proxy_config", mock_proxy_config - ), patch( + with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch( "litellm.proxy.proxy_server.store_model_in_db", True - ), patch( - "litellm.proxy.proxy_server.get_secret_bool", return_value=True - ): + ), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=True): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, prisma_client=mock_prisma_client, @@ -4328,13 +4426,9 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch): mock_proxy_logging.slack_alerting_instance = MagicMock() mock_proxy_config = AsyncMock() - with patch( - "litellm.proxy.proxy_server.proxy_config", mock_proxy_config - ), patch( + with patch("litellm.proxy.proxy_server.proxy_config", mock_proxy_config), patch( "litellm.proxy.proxy_server.store_model_in_db", False - ), patch( - "litellm.proxy.proxy_server.get_secret_bool", return_value=False - ): + ), patch("litellm.proxy.proxy_server.get_secret_bool", return_value=False): # Should not raise an exception await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={},