From 9839b95bc8651c569aea425be2e52696b4899d98 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 31 Jul 2026 12:04:35 -0700 Subject: [PATCH] fix(proxy): seed the applied reload revision at startup A pod adopted whatever revision it found on its first poll, so a manual reload published while the pod was starting was marked applied without ever being served and the pod kept the prices it fetched at import. Read the row once at startup instead, right after that fetch, and treat a missing row as revision 0 --- .../common_utils/periodic_reload_schedule.py | 9 +- litellm/proxy/proxy_server.py | 21 +- .../test_periodic_reload_schedule.py | 6 +- tests/test_litellm/proxy/test_proxy_server.py | 1324 +++++------------ 4 files changed, 412 insertions(+), 948 deletions(-) diff --git a/litellm/proxy/common_utils/periodic_reload_schedule.py b/litellm/proxy/common_utils/periodic_reload_schedule.py index bc353f11ddd..cd0baac4396 100644 --- a/litellm/proxy/common_utils/periodic_reload_schedule.py +++ b/litellm/proxy/common_utils/periodic_reload_schedule.py @@ -137,10 +137,11 @@ def pod_reload_is_due( ) -> bool: """ Whether this pod should reload now. A revision it has not applied means a manual reload - it has not served; ``pod_applied_revision`` is None only until the pod's first poll, when - its data is boot-fresh and it simply adopts whatever revision it finds. Interval reloads - compare against this pod's own data, and a schedule that has never run anywhere fires - immediately rather than one interval later + it has not served; ``pod_applied_revision`` is normally seeded at startup, and is None + only if that read failed, in which case the pod adopts rather than re-serving a request + its boot-time data may already satisfy. Interval reloads compare against this pod's own + data, and a schedule that has never run anywhere fires immediately rather than one + interval later """ if pod_applied_revision is not None and schedule.reload_revision != pod_applied_revision: verbose_proxy_logger.info("%s reload triggered by manual reload request", description) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bf55d39e6ec..b58fc40dd60 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3863,8 +3863,9 @@ class ProxyConfig: ) self.model_cost_map_loaded_at: datetime = get_model_cost_map_loaded_at() or utc_now() - # None until the first poll adopts whatever revision the row carries; this pod's - # data is boot-fresh, so a request published before it started is already satisfied + # Seeded from the row at startup, as close as possible to the boot-time fetch above, + # so a request published while this pod was starting is not mistaken for one its own + # data already satisfies. None only if that read never happened or failed self.model_cost_map_applied_revision: int | None = None def is_yaml(self, config_file_path: str) -> bool: @@ -6439,6 +6440,21 @@ class ProxyConfig: if self._should_load_db_object(object_type="model_cost_map"): await self._check_and_reload_model_cost_map(prisma_client=prisma_client) + async def seed_model_cost_map_revision(self, prisma_client: PrismaClient) -> None: + """ + Adopt the currently published revision once at startup, before the periodic job runs. + + Without this the first poll adopts whatever it finds, so a manual reload published + while this pod was starting would be marked applied without ever being served, and + the pod would keep the prices it fetched at import. A missing row means nobody has + ever requested a reload, which is revision 0. + """ + try: + schedule = await read_reload_schedule(prisma_client, MODEL_COST_MAP_RELOAD_PARAM_NAME) + self.model_cost_map_applied_revision = schedule.reload_revision if schedule is not None else 0 + except Exception as e: + verbose_proxy_logger.exception(f"Error seeding model cost map reload revision: {str(e)}") + async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): """ Check if model cost map needs to be reloaded based on database configuration. @@ -8038,6 +8054,7 @@ class ProxyStartupEvent: config_reload_interval_seconds = 30 ### PERIODIC RELOADS (model cost map, anthropic beta headers) ### + await proxy_config.seed_model_cost_map_revision(prisma_client) scheduler.add_job( proxy_config.check_periodic_reloads, "interval", diff --git a/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py index bef5552f007..eae8c14c54e 100644 --- a/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py +++ b/tests/test_litellm/proxy/common_utils/test_periodic_reload_schedule.py @@ -150,9 +150,9 @@ def test_manual_request_is_identified_not_ordered(): assert (unapplied, applied) == (True, False) -def test_fresh_pod_adopts_the_revision_without_reloading(): - """A pod that just booted holds data newer than any earlier request, so its first poll - must not re-serve one; the interval is the only thing that can make it due""" +def test_pod_without_a_seeded_revision_adopts_without_reloading(): + """Fallback when the startup seed failed: the pod cannot tell whether its boot-time data + already satisfies the request, and re-serving one every restart is the worse guess""" assert ( pod_reload_is_due( schedule=ReloadSchedule(reload_revision=12), diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 325f90e5f44..5fca2e5ad2e 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -19,9 +19,7 @@ from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system-path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system-path import litellm import litellm.proxy.proxy_server as proxy_server_module @@ -178,9 +176,7 @@ def test_login_v2_returns_json_on_http_exception(monkeypatch): from fastapi import HTTPException mock_prisma_client = MagicMock() - mock_authenticate_user = AsyncMock( - side_effect=HTTPException(status_code=401, detail="Unauthorized") - ) + mock_authenticate_user = AsyncMock(side_effect=HTTPException(status_code=401, detail="Unauthorized")) monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", @@ -476,9 +472,7 @@ def test_fallback_login_has_no_deprecation_banner(client_no_auth): "relative/path/logo.png", ], ) -def test_get_logo_url_does_not_disclose_local_paths( - client_no_auth, monkeypatch, ui_logo_path -): +def test_get_logo_url_does_not_disclose_local_paths(client_no_auth, monkeypatch, ui_logo_path): # ``/get_logo_url`` is unauthenticated. Returning a local filesystem # path verbatim discloses admin-only config to any caller. Only # browser-loadable HTTP(S) URLs should be returned; for local paths @@ -578,9 +572,7 @@ def test_restructure_ui_html_files_handles_nested_routes(tmp_path): assert not (ui_root / "home.html").exists() 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" + assert (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" @@ -625,9 +617,7 @@ def test_admin_ui_export_serves_nested_extensionless_routes(): and "_next" not in path.parts and "litellm-asset-prefix" not in path.parts ] - assert not nested_html_offenders, ( - "Nested routes must be named index.html. Offenders: " f"{nested_html_offenders}" - ) + assert not nested_html_offenders, f"Nested routes must be named index.html. Offenders: {nested_html_offenders}" callback_index = out_dir / "mcp" / "oauth" / "callback" / "index.html" assert callback_index.is_file(), ( @@ -644,9 +634,7 @@ def test_admin_ui_export_serves_nested_extensionless_routes(): follow_redirects=False, ) assert redirect.status_code == 307 - assert redirect.headers["location"].endswith( - "/ui/mcp/oauth/callback/?code=abc&state=xyz" - ) + assert redirect.headers["location"].endswith("/ui/mcp/oauth/callback/?code=abc&state=xyz") landed = client.get("/ui/mcp/oauth/callback?code=abc&state=xyz") assert landed.status_code == 200 @@ -749,9 +737,7 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): assert mock_proxy_config.get_credentials.call_count == 1 # Direct call # Verify a scheduled job was added for get_credentials - mock_scheduler_calls = [ - call[0] for call in mock_proxy_config.get_credentials.mock_calls - ] + mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls] assert len(mock_scheduler_calls) > 0 @@ -759,7 +745,9 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypatch): """ Regression (LIT-4882): reload schedules configured from the Admin UI live in the DB and - must fire even without store_model_in_db, which used to gate the job that ran them + must fire even without store_model_in_db, which used to gate the job that ran them. + The pod's revision baseline is seeded before the job can tick, so its first poll cannot + swallow a request published while the pod was starting """ monkeypatch.delenv("DISABLE_PRISMA_SCHEMA_UPDATE", raising=False) monkeypatch.delenv("STORE_MODEL_IN_DB", raising=False) @@ -780,9 +768,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat 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.AsyncIOScheduler", return_value=scheduler - ), + patch("litellm.proxy.proxy_server.AsyncIOScheduler", return_value=scheduler), ): await ProxyStartupEvent.initialize_scheduled_background_jobs( general_settings={}, @@ -795,6 +781,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat assert scheduler.get_job("periodic_reload_job") is not None assert scheduler.get_job("add_deployment_job") is None + mock_proxy_config.seed_model_cost_map_revision.assert_awaited_once_with(mock_prisma_client) finally: scheduler.shutdown(wait=False) @@ -1052,9 +1039,7 @@ def test_get_config_custom_callback_api_env_vars(monkeypatch): assert response.status_code == 200 callbacks = response.json()["callbacks"] - custom_cb = next( - (cb for cb in callbacks if cb["name"] == "custom_callback_api"), None - ) + custom_cb = next((cb for cb in callbacks if cb["name"] == "custom_callback_api"), None) assert custom_cb is not None assert custom_cb["variables"] == { @@ -1102,9 +1087,7 @@ def test_get_config_callbacks_fall_back_to_process_env(mock_env_vars, monkeypatc app.dependency_overrides = original_overrides assert response.status_code == 200 - langfuse_cb = next( - (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None - ) + langfuse_cb = next((cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None) assert langfuse_cb is not None assert langfuse_cb["variables"] == { "LANGFUSE_PUBLIC_KEY": "pk-env-only", @@ -1151,9 +1134,7 @@ def test_get_config_callback_env_secrets_redacted_for_non_admin(mock_env_vars, m app.dependency_overrides = original_overrides assert response.status_code == 200 - langfuse_cb = next( - (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None - ) + langfuse_cb = next((cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None) assert langfuse_cb is not None assert langfuse_cb["variables"]["LANGFUSE_SECRET_KEY"] == "REDACTED" assert langfuse_cb["variables"]["LANGFUSE_HOST"] == "https://cloud.langfuse.com" @@ -1203,9 +1184,7 @@ def test_get_config_returns_email_settings(monkeypatch): app.dependency_overrides = original_overrides assert response.status_code == 200 - email_alert = next( - (a for a in response.json()["alerts"] if a["name"] == "email"), None - ) + email_alert = next((a for a in response.json()["alerts"] if a["name"] == "email"), None) assert email_alert is not None variables = email_alert["variables"] @@ -1350,9 +1329,7 @@ def test_get_config_returns_slack_webhook(monkeypatch): mock_logging = MagicMock() mock_logging.slack_alerting_instance.alert_types = ["budget_alerts"] - mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = [ - "budget_alerts" - ] + mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = ["budget_alerts"] mock_logging.slack_alerting_instance.alert_to_webhook_url = {} monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging) monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) @@ -1369,9 +1346,7 @@ def test_get_config_returns_slack_webhook(monkeypatch): app.dependency_overrides = original_overrides assert response.status_code == 200 - slack_alert = next( - (a for a in response.json()["alerts"] if a["name"] == "slack"), None - ) + slack_alert = next((a for a in response.json()["alerts"] if a["name"] == "slack"), None) assert slack_alert is not None masked_url = slack_alert["variables"]["SLACK_WEBHOOK_URL"] @@ -1391,9 +1366,7 @@ def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): """ from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth - monkeypatch.setenv( - "SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/STALE/OS/ENVVALUE" - ) + monkeypatch.setenv("SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/STALE/OS/ENVVALUE") config_data = { "litellm_settings": {}, "general_settings": {"alerting": ["slack"]}, @@ -1406,9 +1379,7 @@ def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): mock_logging = MagicMock() mock_logging.slack_alerting_instance.alert_types = ["budget_alerts"] - mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = [ - "budget_alerts" - ] + mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = ["budget_alerts"] mock_logging.slack_alerting_instance.alert_to_webhook_url = {} monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging) monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) @@ -1425,9 +1396,7 @@ def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): app.dependency_overrides = original_overrides assert response.status_code == 200 - slack_alert = next( - (a for a in response.json()["alerts"] if a["name"] == "slack"), None - ) + slack_alert = next((a for a in response.json()["alerts"] if a["name"] == "slack"), None) assert slack_alert is not None assert slack_alert["variables"]["SLACK_WEBHOOK_URL"] == "" @@ -1506,9 +1475,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): # Test Case 3: Master key with os.environ prefix test_resolved_key = "sk-resolved-key" - test_config_with_prefix = { - "general_settings": {"master_key": "os.environ/CUSTOM_MASTER_KEY"} - } + test_config_with_prefix = {"general_settings": {"master_key": "os.environ/CUSTOM_MASTER_KEY"}} # Create config with os.environ prefix with open(config_path, "w") as f: @@ -1660,9 +1627,7 @@ async def test_get_all_team_models(): ) # Verify find_many was called with where clause for specific teams - mock_litellm_teamtable.find_many.assert_called_with( - where={"team_id": {"in": ["team1"]}} - ) + mock_litellm_teamtable.find_many.assert_called_with(where={"team_id": {"in": ["team1"]}}) # Verify router.get_model_list was called only for team1 models expected_calls = [ @@ -1857,14 +1822,10 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): prisma_client = MagicMock() prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=2) - prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( - return_value=[db_caller_row, db_other_row] - ) + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[db_caller_row, db_other_row]) caller_user_row = MagicMock() caller_user_row.teams = ["team-mine"] - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=caller_user_row - ) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=caller_user_row) proxy_config = MagicMock() proxy_config.decrypt_model_list_from_db = lambda rows: [ @@ -1894,12 +1855,10 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): assert "byok-db-mine" in filtered_ids assert "public-id" in filtered_ids assert "byok-other" not in filtered_ids, ( - "router-side BYOK from another team must be dropped from search " - "when caller doesn't belong to that team" + "router-side BYOK from another team must be dropped from search when caller doesn't belong to that team" ) assert "byok-db-other" not in filtered_ids, ( - "DB-only BYOK from another team must be dropped from search when " - "caller doesn't belong to that team" + "DB-only BYOK from another team must be dropped from search when caller doesn't belong to that team" ) # total_count is router_models_count (2: caller_team_byok + public_model, # other_team_byok dropped router-side) + DB count (2 from the mocked @@ -2050,9 +2009,7 @@ async def test_filter_models_by_team_id_excludes_viewer_direct_access(): assert "byok-team-111" in visible_ids, "team-111's own BYOK must always be visible" assert "byok-team-222" not in visible_ids, "must not leak other teams' BYOK" - assert ( - "public-id" not in visible_ids - ), "viewer's direct_access must not widen the team's visible set" + assert "public-id" not in visible_ids, "viewer's direct_access must not widen the team's visible set" @pytest.mark.asyncio @@ -2235,9 +2192,7 @@ async def test_add_access_group_models_to_team_models(): mock_ag_row.access_model_names = ["claude-3", "gemini"] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[mock_ag_row] - ) + mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[mock_ag_row]) result = await _add_access_group_models_to_team_models( team_db_objects_typed=[ @@ -2313,9 +2268,7 @@ async def test_add_access_group_models_multiple_teams_shared_group(): mock_extra_row.access_model_names = ["gemini"] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[mock_shared_row, mock_extra_row] - ) + mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[mock_shared_row, mock_extra_row]) result = await _add_access_group_models_to_team_models( team_db_objects_typed=[team_a, team_b], @@ -2507,27 +2460,15 @@ async def test_delete_deployment_type_mismatch(): # The two SHA-hash models have no corresponding entry in combined_id_list # and must be evicted. - assert ( - deleted_count == 2 - ), f"Expected 2 deletions (SHA-hash models), got {deleted_count}" - assert ( - "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695" - in deleted_ids - ) - assert ( - "a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3" - in deleted_ids - ) + assert deleted_count == 2, f"Expected 2 deletions (SHA-hash models), got {deleted_count}" + assert "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695" in deleted_ids + assert "a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3" in deleted_ids # Models 12345678 and 12345679 exist in the config (as integers); str() # conversion in _delete_deployment makes them match the router's string IDs, # so they must NOT be evicted. - assert ( - "12345678" not in deleted_ids - ), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}" - assert ( - "12345679" not in deleted_ids - ), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" + assert "12345678" not in deleted_ids, f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}" + assert "12345679" not in deleted_ids, f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" @pytest.mark.asyncio @@ -2594,9 +2535,7 @@ async def test_get_config_from_file(tmp_path, monkeypatch): await proxy_config._get_config_from_file(str(empty_file)) # Test Case 5: Using global user_config_file_path when no config_file_path provided - monkeypatch.setattr( - "litellm.proxy.proxy_server.user_config_file_path", str(config_file) - ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file)) result = await proxy_config._get_config_from_file(None) assert result == test_config @@ -2715,9 +2654,7 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): ) # 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() @@ -2843,9 +2780,7 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): 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( @@ -2885,9 +2820,7 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp original_max_budget = litellm.max_budget try: proxy_config = ProxyConfig() - await proxy_config.load_config( - router=MagicMock(), config_file_path=str(config_file) - ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) assert isinstance(litellm.max_budget, float) assert litellm.max_budget == 10.0 assert litellm.max_budget > 0 @@ -2922,9 +2855,7 @@ async def test_load_config_max_ui_session_budget_applied_and_coerced(tmp_path, m original_budget = litellm.max_ui_session_budget try: proxy_config = ProxyConfig() - await proxy_config.load_config( - router=MagicMock(), config_file_path=str(config_file) - ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) assert isinstance(litellm.max_ui_session_budget, float) assert litellm.max_ui_session_budget == 2.5 finally: @@ -2950,9 +2881,7 @@ async def test_load_config_max_ui_session_budget_none_disables_cap(tmp_path): original_budget = litellm.max_ui_session_budget try: proxy_config = ProxyConfig() - await proxy_config.load_config( - router=MagicMock(), config_file_path=str(config_file) - ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) assert litellm.max_ui_session_budget is None finally: litellm.max_ui_session_budget = original_budget @@ -3007,10 +2936,7 @@ async def test_load_config_default_internal_user_params_without_max_budget(tmp_p absent_config_file = tmp_path / "absent_config.yaml" absent_config_file.write_text( - "model_list: []\n" - "litellm_settings:\n" - " default_internal_user_params:\n" - " user_role: internal_user\n" + "model_list: []\nlitellm_settings:\n default_internal_user_params:\n user_role: internal_user\n" ) null_config_file = tmp_path / "null_config.yaml" @@ -3057,9 +2983,7 @@ async def test_load_config_user_url_validation_handles_null_and_string_false(tmp ) ) - await ProxyConfig().load_config( - router=MagicMock(), config_file_path=str(null_config_file) - ) + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(null_config_file)) assert litellm.user_url_validation is True assert litellm.user_url_allowed_hosts is None assert litellm.provider_url_destination_allowed_hosts is None @@ -3074,9 +2998,7 @@ async def test_load_config_user_url_validation_handles_null_and_string_false(tmp ) ) - await ProxyConfig().load_config( - router=MagicMock(), config_file_path=str(false_config_file) - ) + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(false_config_file)) assert litellm.user_url_validation is False @@ -3104,12 +3026,8 @@ async def test_load_environment_variables_direct_and_os_environ(): # 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) @@ -3122,9 +3040,7 @@ async def test_load_environment_variables_direct_and_os_environ(): 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 @@ -3177,9 +3093,7 @@ async def test_load_environment_variables_litellm_license_and_edge_cases(): 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"} - } + test_config_secret_none = {"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): @@ -3218,9 +3132,7 @@ async def test_load_environment_variables_blocks_dangerous_keys(): # 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 "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 @@ -3294,15 +3206,11 @@ async def test_write_config_to_file(monkeypatch): # Mock general_settings mock_general_settings = {"store_model_in_db": True} - monkeypatch.setattr( - "litellm.proxy.proxy_server.general_settings", mock_general_settings - ) + 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 - ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", test_config_path) proxy_config = ProxyConfig() @@ -3323,9 +3231,7 @@ async def test_write_config_to_file(monkeypatch): # 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["data"] == {"key": "value"} # model_list should be popped assert call_args.kwargs["table_name"] == "config" @@ -3346,15 +3252,11 @@ async def test_write_config_to_file_when_store_model_in_db_false(monkeypatch): # Mock general_settings mock_general_settings = {"store_model_in_db": False} - monkeypatch.setattr( - "litellm.proxy.proxy_server.general_settings", mock_general_settings - ) + 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 - ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", test_config_path) proxy_config = ProxyConfig() @@ -3409,22 +3311,20 @@ async def test_async_data_generator_midstream_error(): for chunk in mock_chunks: yield chunk - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator # Mock async_post_call_streaming_hook to return error on third chunk def mock_streaming_hook(*args, **kwargs): chunk = kwargs.get("response") # Return error message for the third chunk (simulating guardrail trigger) if chunk == mock_chunks[2]: - return 'data: {"error": {"error": "Azure Content Safety Guardrail: Hate crossed severity 2, Got severity: 2"}}' + return ( + 'data: {"error": {"error": "Azure Content Safety Guardrail: Hate crossed severity 2, Got severity: 2"}}' + ) # Return normal chunks for first two return chunk - mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( - side_effect=mock_streaming_hook - ) + mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=mock_streaming_hook) mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() # Mock the global proxy_logging_obj @@ -3435,26 +3335,18 @@ async def test_async_data_generator_midstream_error(): # Collect all yielded data from the generator yielded_data = [] try: - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) except Exception as e: # If there's an exception, that's also part of what we want to test pass # Verify the results - assert ( - len(yielded_data) >= 3 - ), f"Expected at least 3 chunks, got {len(yielded_data)}: {yielded_data}" + assert len(yielded_data) >= 3, f"Expected at least 3 chunks, got {len(yielded_data)}: {yielded_data}" # First two chunks should be normal data - assert yielded_data[0].startswith( - "data: " - ), f"First chunk should start with 'data: ', got: {yielded_data[0]}" - assert yielded_data[1].startswith( - "data: " - ), f"Second chunk should start with 'data: ', got: {yielded_data[1]}" + assert yielded_data[0].startswith("data: "), f"First chunk should start with 'data: ', got: {yielded_data[0]}" + assert yielded_data[1].startswith("data: "), f"Second chunk should start with 'data: ', got: {yielded_data[1]}" # The error message should be yielded error_found = False @@ -3466,15 +3358,11 @@ async def test_async_data_generator_midstream_error(): if "data: [DONE]" in data: done_found = True - assert ( - error_found - ), f"Error message should be found in yielded data. Got: {yielded_data}" + assert error_found, f"Error message should be found in yielded data. Got: {yielded_data}" assert done_found, f"[DONE] message should be found at the end. Got: {yielded_data}" # Verify that the streaming hook was called for each chunk - assert mock_proxy_logging_obj.async_post_call_streaming_hook.call_count == len( - mock_chunks - ) + assert mock_proxy_logging_obj.async_post_call_streaming_hook.call_count == len(mock_chunks) # Verify that post_call_failure_hook was NOT called (since this is not an exception case) mock_proxy_logging_obj.post_call_failure_hook.assert_not_called() @@ -3561,15 +3449,11 @@ async def test_chat_completion_result_no_nested_none_values(): # Verify the mock has None values before serialization raw_dict = mock_model_response.model_dump() none_paths_before = _has_nested_none_values(raw_dict) - assert ( - len(none_paths_before) > 0 - ), "Mock should have None values before exclude_none=True" + assert len(none_paths_before) > 0, "Mock should have None values before exclude_none=True" # Mock the request processing to return our mock response mock_base_processor = MagicMock() - mock_base_processor.base_process_llm_request = AsyncMock( - return_value=mock_model_response - ) + mock_base_processor.base_process_llm_request = AsyncMock(return_value=mock_model_response) # Mock other dependencies mock_request = MagicMock(spec=Request) @@ -3598,9 +3482,9 @@ async def test_chat_completion_result_no_nested_none_values(): # Check that there are no nested None values in the result none_paths_after = _has_nested_none_values(result) - assert ( - len(none_paths_after) == 0 - ), f"Result should not contain nested None values. Found None at: {none_paths_after}" + assert len(none_paths_after) == 0, ( + f"Result should not contain nested None values. Found None at: {none_paths_after}" + ) # Verify essential fields are present assert "id" in result @@ -3626,9 +3510,7 @@ async def test_chat_completion_result_no_nested_none_values(): "annotations", ] for field in excluded_fields: - assert ( - field not in message - ), f"Field '{field}' should be excluded when it's None" + assert field not in message, f"Field '{field}' should be excluded when it's None" # ============================================================================ @@ -3678,12 +3560,8 @@ class TestPriceDataReloadAPI: # subsequent tests running in the same worker process. original_model_cost = litellm.model_cost.copy() try: - with patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map: - mock_get_map.return_value = { - "gpt-3.5-turbo": {"input_cost_per_token": 0.001} - } + with patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map: + mock_get_map.return_value = {"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: mock_prisma.db.litellm_config.upsert = AsyncMock( @@ -3699,10 +3577,7 @@ class TestPriceDataReloadAPI: assert "timestamp" in data assert "models_count" in data # The new implementation immediately reloads and returns the count - assert ( - "Price data reloaded successfully! 1 models updated." - in data["message"] - ) + assert "Price data reloaded successfully! 1 models updated." in data["message"] assert data["models_count"] == 1 finally: # Restore the full model cost map so subsequent tests are not affected @@ -3725,9 +3600,7 @@ class TestPriceDataReloadAPI: def test_get_model_cost_map_public_access(self, client_no_auth): """Test that the model cost map endpoint is publicly accessible""" - with patch( - "litellm.model_cost", {"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} - ): + with patch("litellm.model_cost", {"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}): response = client_no_auth.get("/public/litellm_model_cost_map") assert response.status_code == 200 @@ -3736,23 +3609,19 @@ class TestPriceDataReloadAPI: def test_reload_model_cost_map_error_handling(self, client_with_auth): """Test error handling in the reload endpoint""" - with patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map: + with patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map: mock_get_map.side_effect = Exception("Network error") # 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.upsert = AsyncMock( - return_value=_reload_schedule_row({}, reload_revision=1) - ) + return_value=_reload_schedule_row({}, reload_revision=1) + ) response = client_with_auth.post("/reload/model_cost_map") - assert ( - response.status_code == 500 - ) # The new implementation immediately reloads and fails on error + assert response.status_code == 500 # The new implementation immediately reloads and fails on error data = response.json() assert "Failed to reload model cost map" in data["detail"] @@ -3760,9 +3629,7 @@ class TestPriceDataReloadAPI: """Admin schedule write owns param_value only, so it can't clobber the job-owned run columns""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: # Mock database upsert - mock_prisma.db.litellm_config.upsert = AsyncMock( - return_value=_reload_schedule_row({}, reload_revision=1) - ) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1)) response = client_with_auth.post("/schedule/model_cost_map_reload?hours=6") @@ -3774,9 +3641,7 @@ class TestPriceDataReloadAPI: assert "timestamp" in data call_args = mock_prisma.db.litellm_config.upsert.call_args - assert call_args[1]["where"] == { - "param_name": "model_cost_map_reload_config" - } + assert call_args[1]["where"] == {"param_name": "model_cost_map_reload_config"} update_payload = call_args[1]["data"]["update"] assert set(update_payload.keys()) == {"param_value"} assert json.loads(update_payload["param_value"]) == {"interval_hours": 6} @@ -3839,9 +3704,7 @@ class TestPriceDataReloadAPI: Regression (LIT-4882): status is served purely from the DB row, so a restarted pod (whose in-memory clock only knows its own boot) still reports the real last/next run """ - proxy_server_module.proxy_config.model_cost_map_loaded_at = datetime( - 2030, 6, 1, tzinfo=timezone.utc - ) + proxy_server_module.proxy_config.model_cost_map_loaded_at = datetime(2030, 6, 1, tzinfo=timezone.utc) with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: mock_prisma.db.litellm_config.find_unique = AsyncMock( @@ -3958,9 +3821,7 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: - with patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map: + with patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map: mock_get_map.return_value = mock_cost_map # Mock the database connection @@ -4003,9 +3864,7 @@ class TestPriceDataReloadIntegration: proxy_config = ProxyConfig() mock_prisma = MagicMock() - mock_prisma.db.litellm_config.upsert = AsyncMock( - return_value=_reload_schedule_row({}, reload_revision=1) - ) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1)) mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) @@ -4029,20 +3888,14 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map, + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = { - "gpt-3.5-turbo": {"input_cost_per_token": 0.001} - } + mock_get_map.return_value = {"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) - assert litellm.model_cost["gpt-3.5-turbo"] == { - "input_cost_per_token": 0.001 - } + assert litellm.model_cost["gpt-3.5-turbo"] == {"input_cost_per_token": 0.001} assert proxy_config.model_cost_map_loaded_at == frozen_now assert mock_prisma.db.litellm_config.update_many.call_args[1] == { "data": {"last_run_at": frozen_now}, @@ -4079,9 +3932,7 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map, + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4111,16 +3962,12 @@ class TestPriceDataReloadIntegration: ) mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) - proxy_config.model_cost_map_loaded_at = datetime( - 2024, 1, 1, 0, 0, tzinfo=timezone.utc - ) + proxy_config.model_cost_map_loaded_at = datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc) original_model_cost = litellm.model_cost.copy() try: with ( - patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map, + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = {"gpt-4-test": {"input_cost_per_token": 0.5}} @@ -4129,9 +3976,7 @@ class TestPriceDataReloadIntegration: assert litellm.model_cost["gpt-4-test"] == {"input_cost_per_token": 0.5} assert proxy_config.model_cost_map_loaded_at == frozen_now - assert mock_prisma.db.litellm_config.update_many.call_args[1][ - "data" - ] == {"last_run_at": frozen_now} + assert mock_prisma.db.litellm_config.update_many.call_args[1]["data"] == {"last_run_at": frozen_now} mock_get_map.reset_mock() mock_prisma.db.litellm_config.update_many.reset_mock() @@ -4141,9 +3986,7 @@ class TestPriceDataReloadIntegration: mock_get_map.assert_not_called() mock_prisma.db.litellm_config.update_many.assert_not_called() - assert proxy_config.model_cost_map_loaded_at == frozen_now - timedelta( - hours=1 - ) + assert proxy_config.model_cost_map_loaded_at == frozen_now - timedelta(hours=1) finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() @@ -4157,9 +4000,7 @@ class TestPriceDataReloadIntegration: frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) mock_prisma = MagicMock() mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) - mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=_reload_schedule_row({}, reload_revision=1) - ) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1)) for pod in pods: pod.model_cost_map_applied_revision = 0 pod.model_cost_map_loaded_at = frozen_now @@ -4167,9 +4008,7 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map, + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} @@ -4184,6 +4023,67 @@ class TestPriceDataReloadIntegration: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() + def test_startup_seed_lets_a_boot_race_request_through(self): + """ + Regression: a manual reload published after this pod started must still be served. + Seeding the applied revision at startup is the only thing that distinguishes it + from a request the pod's own boot-time fetch already satisfied + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + proxy_config.model_cost_map_loaded_at = frozen_now + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=4)) + + asyncio.run(proxy_config.seed_model_cost_map_revision(mock_prisma)) + assert proxy_config.model_cost_map_applied_revision == 4 + + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=5)) + original_model_cost = litellm.model_cost.copy() + try: + with ( + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + ): + mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} + + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + mock_get_map.assert_called_once() + assert proxy_config.model_cost_map_applied_revision == 5 + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + def test_startup_seed_treats_a_missing_row_as_never_requested(self): + """No row means nobody has ever asked for a reload, so the seed is 0; leaving it + unset would make the very first request adopted-but-never-served on this pod""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=None) + + asyncio.run(proxy_config.seed_model_cost_map_revision(mock_prisma)) + + assert proxy_config.model_cost_map_applied_revision == 0 + + def test_startup_seed_failure_does_not_block_startup(self): + """A database hiccup at boot must not take the proxy down; the pod falls back to + adopting on its first poll, which is the pre-seed behavior""" + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock(side_effect=Exception("connection refused")) + + asyncio.run(proxy_config.seed_model_cost_map_revision(mock_prisma)) + + assert proxy_config.model_cost_map_applied_revision is None + def test_distributed_reload_stamps_last_run_without_creating_row(self): """ Regression: the job's write carries neither param_value (which would clobber the @@ -4194,21 +4094,15 @@ class TestPriceDataReloadIntegration: proxy_config = ProxyConfig() mock_prisma = MagicMock() - mock_prisma.db.litellm_config.find_unique = AsyncMock( - return_value=_reload_schedule_row({"interval_hours": 24}) - ) - mock_prisma.db.litellm_config.upsert = AsyncMock( - return_value=_reload_schedule_row({}, reload_revision=1) - ) + mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=_reload_schedule_row({"interval_hours": 24})) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1)) mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) original_model_cost = litellm.model_cost.copy() try: with ( - patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map, + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): mock_get_map.return_value = {"gpt-4": {"input_cost_per_token": 0.001}} @@ -4247,9 +4141,7 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch( - "litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map" - ) as mock_get_map, + patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map, patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): @@ -4272,13 +4164,10 @@ class TestPriceDataReloadIntegration: "last_run_at": frozen_now, "reload_revision": 1, } - assert ( - proxy_server_module.proxy_config.model_cost_map_loaded_at - == frozen_now + assert proxy_server_module.proxy_config.model_cost_map_loaded_at == frozen_now + assert proxy_server_module.proxy_config.model_cost_map_applied_revision == 9, ( + "the serving pod must adopt the revision it published, not reload again" ) - assert ( - proxy_server_module.proxy_config.model_cost_map_applied_revision == 9 - ), "the serving pod must adopt the revision it published, not reload again" finally: litellm.model_cost = original_model_cost _invalidate_model_cost_lowercase_map() @@ -4302,18 +4191,12 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.find_unique = AsyncMock(return_value=mock_config) # _check_and_reload_anthropic_beta_headers now reads through get_generic_data. mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) - mock_prisma.db.litellm_config.upsert = AsyncMock( - return_value=_reload_schedule_row({}, reload_revision=1) - ) + mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1)) - with patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload: + with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") 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() @@ -4345,21 +4228,17 @@ class TestPriceDataReloadIntegration: app.dependency_overrides[user_api_key_auth] = lambda: mock_auth client = TestClient(app) - with patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload: + with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") as mock_reload: mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} 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": 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=_reload_schedule_row({}, reload_revision=1) - ) + return_value=_reload_schedule_row({}, reload_revision=1) + ) response = client.post("/reload/anthropic_beta_headers") assert response.status_code == 200 @@ -4447,9 +4326,7 @@ async def test_add_router_settings_from_db_config_merge_logic(): # Mock prisma client mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) # Call the method under test await proxy_config._add_router_settings_from_db_config( @@ -4459,9 +4336,7 @@ async def test_add_router_settings_from_db_config_merge_logic(): ) # Verify find_first was called with correct parameters - mock_prisma_client.db.litellm_config.find_first.assert_called_once_with( - where={"param_name": "router_settings"} - ) + mock_prisma_client.db.litellm_config.find_first.assert_called_once_with(where={"param_name": "router_settings"}) # Verify update_settings was called mock_router.update_settings.assert_called_once() @@ -4541,9 +4416,7 @@ async def test_add_router_settings_from_db_config_edge_cases(): # Test Case 4: Config has no router_settings mock_db_config = MagicMock() mock_db_config.param_value = {"db_setting": "db_value"} - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) await proxy_config._add_router_settings_from_db_config( config_data={}, # No router_settings in config @@ -4568,9 +4441,7 @@ async def test_add_router_settings_from_db_config_edge_cases(): # Test Case 6: DB config exists but param_value is not a dict mock_db_config_invalid = MagicMock() mock_db_config_invalid.param_value = "not_a_dict" - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config_invalid - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config_invalid) config_data = {"router_settings": {"config_setting": "config_value"}} @@ -4622,9 +4493,7 @@ async def test_add_router_settings_shallow_merge_behavior(): } mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) await proxy_config._add_router_settings_from_db_config( config_data=config_data, @@ -4701,9 +4570,7 @@ async def test_model_info_v1_oci_secrets_not_leaked(): 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 - ) + result = await model_info_v1(user_api_key_dict=mock_user_api_key_dict, litellm_model_id=None) # Verify the result structure assert "data" in result @@ -4714,40 +4581,24 @@ async def test_model_info_v1_oci_secrets_not_leaked(): # 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 @@ -4777,9 +4628,7 @@ 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 @@ -4789,9 +4638,7 @@ 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 @@ -4892,10 +4739,7 @@ def test_should_load_db_object_with_supported_db_objects(): assert proxy_config._should_load_db_object(object_type="mcp") is True assert proxy_config._should_load_db_object(object_type="guardrails") is True assert proxy_config._should_load_db_object(object_type="vector_stores") is True - assert ( - proxy_config._should_load_db_object(object_type="pass_through_endpoints") - is True - ) + assert proxy_config._should_load_db_object(object_type="pass_through_endpoints") is True assert proxy_config._should_load_db_object(object_type="prompts") is True assert proxy_config._should_load_db_object(object_type="model_cost_map") is True @@ -4921,12 +4765,8 @@ 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, @@ -4980,9 +4820,7 @@ async def test_tag_cache_update_multiple_tags(): 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_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -5003,9 +4841,7 @@ 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 @@ -5031,9 +4867,7 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl(): "async_get_cache", new=AsyncMock(return_value={"tag_name": "active-tag", "spend": 1.0}), ): - with patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_set_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, @@ -5076,9 +4910,7 @@ async def test_spend_tracking_never_writes_the_auth_object_back(): model_type=UserAPIKeyAuth, ) with ( - patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_pipeline, + patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_pipeline, patch.object(cache, "async_set_cache", new=AsyncMock()) as mock_set, ): await litellm.proxy.proxy_server.update_cache( @@ -5089,9 +4921,7 @@ async def test_spend_tracking_never_writes_the_auth_object_back(): response_cost=5.0, parent_otel_span=None, ) - pending = [ - t for t in asyncio.all_tasks() if t is not asyncio.current_task() - ] + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] if pending: await asyncio.wait(pending, timeout=5) @@ -5133,12 +4963,8 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared(): cache = DualCache(default_in_memory_ttl=300) setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) try: - with patch.object( - cache, "async_get_cache", new=AsyncMock(side_effect=fake_get) - ): - 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=fake_get)): + 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="user-lit", @@ -5148,24 +4974,14 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared(): parent_otel_span=None, ) - pending = [ - t for t in asyncio.all_tasks() if t is not asyncio.current_task() - ] + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] if pending: await asyncio.wait(pending, timeout=5) calls = mock_set_cache.await_args_list - local_keys = [ - k - for c in calls - if c.kwargs.get("local_only") is True - for k, _ in c.kwargs["cache_list"] - ] + local_keys = [k for c in calls if c.kwargs.get("local_only") is True for k, _ in c.kwargs["cache_list"]] shared_keys = [ - k - for c in calls - if c.kwargs.get("local_only") is not True - for k, _ in c.kwargs["cache_list"] + k for c in calls if c.kwargs.get("local_only") is not True for k, _ in c.kwargs["cache_list"] ] assert "user-lit" in local_keys assert global_key not in local_keys @@ -5196,20 +5012,14 @@ async def test_init_sso_settings_in_db(): } mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - return_value=mock_sso_config - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_sso_config) # Mock _decrypt_and_set_db_env_variables - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt_and_set: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt_and_set: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) # Verify find_unique was called with correct parameters - mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( - where={"id": "sso_config"} - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(where={"id": "sso_config"}) # Verify _decrypt_and_set_db_env_variables was called with uppercased keys mock_decrypt_and_set.assert_called_once() @@ -5249,15 +5059,11 @@ async def test_init_sso_settings_in_db_no_settings(): mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) # Mock _decrypt_and_set_db_env_variables - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt_and_set: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt_and_set: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) # Verify find_unique was called - mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( - where={"id": "sso_config"} - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(where={"id": "sso_config"}) # Verify _decrypt_and_set_db_env_variables was NOT called when no settings exist mock_decrypt_and_set.assert_not_called() @@ -5276,9 +5082,7 @@ async def test_init_sso_settings_in_db_error_handling(): # Mock prisma client to raise an exception mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - side_effect=Exception("Database connection error") - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(side_effect=Exception("Database connection error")) # The method should not raise an exception, it should log it instead try: @@ -5287,9 +5091,7 @@ 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 @@ -5308,20 +5110,14 @@ async def test_init_sso_settings_in_db_empty_settings(): mock_sso_config.sso_settings = {} mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - return_value=mock_sso_config - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_sso_config) # Mock _decrypt_and_set_db_env_variables - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt_and_set: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt_and_set: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) # Verify find_unique was called - mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( - where={"id": "sso_config"} - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(where={"id": "sso_config"}) # Verify _decrypt_and_set_db_env_variables was called with empty dict mock_decrypt_and_set.assert_called_once() @@ -5354,16 +5150,12 @@ async def test_init_sso_settings_in_db_retries_on_transport_error(): return mock_sso_config mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - side_effect=_flaky_find_unique - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(side_effect=_flaky_find_unique) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) assert len(invocations) == 2 @@ -5384,9 +5176,7 @@ async def test_init_sso_settings_in_db_propagates_when_reconnect_fails(): proxy_config = ProxyConfig() mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - side_effect=prisma.errors.ClientNotConnectedError() - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(side_effect=prisma.errors.ClientNotConnectedError()) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -5417,24 +5207,17 @@ async def test_init_hashicorp_vault_config_override_retries_on_transport_error() return None # No config in DB → function returns early after retry. mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock( - side_effect=_flaky_find_unique - ) + mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock(side_effect=_flaky_find_unique) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 - await proxy_config._init_hashicorp_vault_config_override( - prisma_client=mock_prisma_client - ) + await proxy_config._init_hashicorp_vault_config_override(prisma_client=mock_prisma_client) assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "init_hashicorp_vault_config_override_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "init_hashicorp_vault_config_override_lookup_failure" def test_update_config_fields_uppercases_env_vars(monkeypatch): @@ -5484,37 +5267,20 @@ def test_encrypt_env_variables_for_db_is_idempotent(monkeypatch): plaintext = "pk-langfuse-secret-value" # First write: plaintext in -> single-encrypted out. - enc1 = proxy_config._encrypt_env_variables_for_db( - {"LANGFUSE_PUBLIC_KEY": plaintext} - ) + enc1 = proxy_config._encrypt_env_variables_for_db({"LANGFUSE_PUBLIC_KEY": plaintext}) assert enc1["LANGFUSE_PUBLIC_KEY"] != plaintext - assert ( - decrypt_value_helper( - value=enc1["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY" - ) - == plaintext - ) + assert decrypt_value_helper(value=enc1["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY") == plaintext # UI round-trip: feed the ciphertext back in. Must NOT double-encrypt. enc2 = proxy_config._encrypt_env_variables_for_db(enc1) - assert ( - decrypt_value_helper( - value=enc2["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY" - ) - == plaintext - ) + assert decrypt_value_helper(value=enc2["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY") == plaintext # And again, ×3 total ciphertext re-feeds — still exactly one layer, # never stacked, no matter how many times the UI re-saves. enc3 = proxy_config._encrypt_env_variables_for_db(enc2) enc4 = proxy_config._encrypt_env_variables_for_db(enc3) for stacked in (enc3, enc4): - assert ( - decrypt_value_helper( - value=stacked["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY" - ) - == plaintext - ) + assert decrypt_value_helper(value=stacked["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY") == plaintext # Write path must not leak the value into the process environment. assert os.environ.get("LANGFUSE_PUBLIC_KEY") is None @@ -5556,15 +5322,11 @@ 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" @@ -5632,9 +5394,7 @@ async def test_get_image_non_root_uses_var_lib_assets_dir(monkeypatch): 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.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, @@ -5684,9 +5444,7 @@ async def test_get_image_non_root_fallback_to_default_logo(monkeypatch): # 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.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, @@ -5709,9 +5467,7 @@ 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" @@ -5751,14 +5507,8 @@ async def test_get_image_root_case_uses_current_dir(monkeypatch): await get_image() # 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 - 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" + var_lib_assets_calls = [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" # Verify FileResponse was called assert mock_file_response.called, "FileResponse should be called" @@ -5789,15 +5539,11 @@ async def test_get_image_custom_local_logo_bypasses_cache(monkeypatch, tmp_path) return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + 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] == str(custom_logo.resolve()), ( f"Expected custom logo path, got {calls_to_file_response[0]}. " "A stale cached_logo.jpg may have been returned instead." @@ -5827,24 +5573,18 @@ async def test_get_image_default_logo_ignores_stale_cache(monkeypatch, tmp_path) return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + 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 != str(cache_path.resolve()) assert served_path.endswith("logo.jpg") @pytest.mark.asyncio -async def test_get_image_custom_logo_missing_falls_through_to_default( - monkeypatch, tmp_path -): +async def test_get_image_custom_logo_missing_falls_through_to_default(monkeypatch, tmp_path): """ Test that when UI_LOGO_PATH points to a non-existent local file, get_image falls through to the default logo instead of failing. @@ -5865,26 +5605,18 @@ async def test_get_image_custom_logo_missing_falls_through_to_default( return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + 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 != str( - custom_logo_path - ), "Should not attempt to serve a non-existent custom logo" + assert served_path != str(custom_logo_path), "Should not attempt to serve a non-existent custom logo" assert served_path.endswith("logo.jpg") @pytest.mark.asyncio -async def test_get_image_custom_logo_missing_no_cache_serves_default( - monkeypatch, tmp_path -): +async def test_get_image_custom_logo_missing_no_cache_serves_default(monkeypatch, tmp_path): """ Test that when UI_LOGO_PATH points to a non-existent file AND there is no cached_logo.jpg, get_image serves the default logo instead of the non-existent @@ -5906,22 +5638,14 @@ async def test_get_image_custom_logo_missing_no_cache_serves_default( return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + 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 != str( - custom_logo_path - ), "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 != str(custom_logo_path), "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): @@ -5961,9 +5685,7 @@ def test_get_config_normalizes_string_callbacks(monkeypatch): success_callbacks = [cb["name"] for cb in callbacks if cb.get("type") == "success"] failure_callbacks = [cb["name"] for cb in callbacks if cb.get("type") == "failure"] - success_and_failure_callbacks = [ - cb["name"] for cb in callbacks if cb.get("type") == "success_and_failure" - ] + success_and_failure_callbacks = [cb["name"] for cb in callbacks if cb.get("type") == "success_and_failure"] assert "langfuse" in success_callbacks assert len(failure_callbacks) == 0 @@ -6000,9 +5722,7 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): }, } - result = proxy_config._update_config_fields( - current_config, "general_settings", db_param_value - ) + result = proxy_config._update_config_fields(current_config, "general_settings", db_param_value) assert result["general_settings"]["max_parallel_requests"] == 10 assert result["general_settings"]["allowed_models"] == ["gpt-3.5-turbo", "gpt-4"] @@ -6069,9 +5789,7 @@ class TestInvitationEndpoints: ), ], ) - def test_invitation_endpoints_proxy_admin_success( - self, client_with_auth, endpoint, payload, mock_return - ): + def test_invitation_endpoints_proxy_admin_success(self, client_with_auth, endpoint, payload, mock_return): """Proxy admin can successfully create and delete invitations.""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: mock_prisma.db.litellm_invitationlink = MagicMock() @@ -6086,9 +5804,7 @@ class TestInvitationEndpoints: mock_prisma.db.litellm_invitationlink.find_unique = AsyncMock( return_value={**mock_return, "created_by": "admin-user-id"} ) - mock_prisma.db.litellm_invitationlink.delete = AsyncMock( - return_value=mock_return - ) + mock_prisma.db.litellm_invitationlink.delete = AsyncMock(return_value=mock_return) response = client_with_auth.post(endpoint, json=payload) assert response.status_code == 200 @@ -6103,9 +5819,7 @@ class TestInvitationEndpoints: ("/invitation/delete", {"invitation_id": "inv-456"}), ], ) - def test_invitation_endpoints_non_admin_denied( - self, client_with_auth, endpoint, payload - ): + def test_invitation_endpoints_non_admin_denied(self, client_with_auth, endpoint, payload): """Non-admin users cannot access invitation endpoints.""" from litellm.proxy._types import LitellmUserRoles @@ -6160,9 +5874,7 @@ async def test_async_data_generator_cleanup_on_early_exit(): for chunk in mock_chunks: yield chunk - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs.get("response") ) @@ -6174,9 +5886,7 @@ async def test_async_data_generator_cleanup_on_early_exit(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): # Consume only the first chunk then abandon the generator (simulates client disconnect) - gen = async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ) + gen = async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data) first_chunk = await gen.__anext__() assert first_chunk.startswith("data: ") @@ -6229,19 +5939,12 @@ async def test_async_data_generator_uses_direct_stream_fast_path_without_callbac mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): - with patch.object( - ProxyLogging, "_fire_deferred_stream_logging" - ) as mock_deferred_logging: + with patch.object(ProxyLogging, "_fire_deferred_stream_logging") as mock_deferred_logging: yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert len([chunk for chunk in yielded_text if chunk.startswith("data: {")]) == 2 assert yielded_text[-1] == "data: [DONE]\n\n" mock_proxy_logging_obj.async_post_call_streaming_iterator_hook.assert_not_called() @@ -6294,18 +5997,13 @@ async def test_async_data_generator_preserves_non_raw_sse_like_bytes(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text[0] == gemini_event.decode("utf-8") assert yielded_text[1] == gemini_event_without_terminator.decode("utf-8") + "\n\n" - assert yielded_text[2] == f'data: {raw_payload.decode("utf-8")}\n\n' + assert yielded_text[2] == f"data: {raw_payload.decode('utf-8')}\n\n" assert "b'data:" not in "".join(yielded_text) assert yielded_text[-1] == "data: [DONE]\n\n" @@ -6328,12 +6026,8 @@ async def test_async_data_generator_buffers_split_google_native_sse_json_frame() ) raw_chunks = [ payload[:2].encode("utf-8"), - payload[ - 2 : payload.index("thoughtSignature") + len('thoughtSignature": "abc') - ].encode("utf-8"), - payload[ - payload.index("thoughtSignature") + len('thoughtSignature": "abc') : - ].encode("utf-8"), + payload[2 : payload.index("thoughtSignature") + len('thoughtSignature": "abc')].encode("utf-8"), + payload[payload.index("thoughtSignature") + len('thoughtSignature": "abc') :].encode("utf-8"), ] class MockStream: @@ -6360,15 +6054,10 @@ async def test_async_data_generator_buffers_split_google_native_sse_json_frame() with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text == [payload] for chunk in yielded_text: @@ -6414,15 +6103,10 @@ async def test_async_data_generator_flushes_raw_sse_stream_without_trailing_deli patch.object(ProxyLogging, "_fire_deferred_stream_logging"), ): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert len(yielded_text) == 1 assert yielded_text[0] == 'data: {"candidates": [{"content": "unterminated"}]\n\n' assert "[DONE]" not in yielded_text[0] @@ -6469,15 +6153,10 @@ async def test_async_data_generator_errors_when_raw_sse_frame_exceeds_buffer_lim patch.object(ProxyLogging, "_fire_deferred_stream_logging"), ): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert len(yielded_text) == 1 assert "maximum buffered size" in yielded_text[0] assert "[DONE]" not in yielded_text[0] @@ -6530,15 +6209,10 @@ async def test_async_data_generator_checks_raw_sse_buffer_limit_after_complete_f patch.object(ProxyLogging, "_fire_deferred_stream_logging"), ): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text[0] == complete_frame assert yielded_text[1] == partial_frame + "\n\n" assert "[DONE]" not in "".join(yielded_text) @@ -6559,9 +6233,7 @@ async def test_async_data_generator_google_genai_stream_omits_openai_done(): "model": "gemini-2.0-flash", "_litellm_skip_openai_stream_done": True, } - gemini_event = ( - b'data: {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]}\n\n' - ) + gemini_event = b'data: {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]}\n\n' class MockStream: def __aiter__(self): @@ -6586,15 +6258,10 @@ async def test_async_data_generator_google_genai_stream_omits_openai_done(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text == [gemini_event.decode("utf-8")] assert "[DONE]" not in "".join(yielded_text) @@ -6683,15 +6350,10 @@ async def test_async_data_generator_google_genai_stream_forwards_error_without_d with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text == [error_sse] assert "[DONE]" not in "".join(yielded_text) @@ -6721,9 +6383,7 @@ async def test_async_data_generator_cleanup_on_normal_completion(): for chunk in mock_chunks: yield chunk - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs.get("response") ) @@ -6734,9 +6394,7 @@ async def test_async_data_generator_cleanup_on_normal_completion(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) # Should have completed normally with [DONE] @@ -6767,9 +6425,7 @@ async def test_async_data_generator_cleanup_on_midstream_error(): yield {"choices": [{"delta": {"content": "Hello"}}]} raise RuntimeError("upstream connection reset") - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator_with_error - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator_with_error mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs.get("response") ) @@ -6780,9 +6436,7 @@ async def test_async_data_generator_cleanup_on_midstream_error(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) # Should have yielded data chunk and then an error chunk @@ -6837,9 +6491,7 @@ async def test_update_general_settings_store_model_in_db_true(): patch("litellm.proxy.proxy_server.store_model_in_db", False) as mock_store, patch("litellm.proxy.proxy_server.general_settings", {}) as mock_gs, ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": True} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True}) import litellm.proxy.proxy_server as ps @@ -6861,9 +6513,7 @@ async def test_update_general_settings_store_model_in_db_false(): 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} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": False}) import litellm.proxy.proxy_server as ps @@ -6911,9 +6561,7 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): 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"} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": "true"}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is True @@ -6923,9 +6571,7 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): 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"} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": "True"}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is True @@ -6935,9 +6581,7 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): 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"} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": "false"}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is False @@ -6958,9 +6602,7 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): 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} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": None}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is True @@ -6970,9 +6612,7 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): 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} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": None}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is False @@ -6992,9 +6632,7 @@ async def test_store_model_in_db_db_override_when_config_false(): # Mock DB returning store_model_in_db=True in general_settings mock_db_record = MagicMock() mock_db_record.param_value = {"store_model_in_db": True} - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_record - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_record) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() @@ -7078,9 +6716,7 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch): mock_prisma_client = MagicMock() # Simulate DB failure - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - side_effect=Exception("DB connection error") - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(side_effect=Exception("DB connection error")) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() @@ -7218,9 +6854,7 @@ async def test_increment_spend_counters_initializes_and_increments(): ) # Counter should be: base(5.0) + increment(0.50) = 5.50 - counter = counter_cache.in_memory_cache.get_cache( - key=f"spend:key:{hashed_token}" - ) + counter = counter_cache.in_memory_cache.get_cache(key=f"spend:key:{hashed_token}") assert counter == 5.50 # Second increment — counter already exists, just increment @@ -7231,9 +6865,7 @@ async def test_increment_spend_counters_initializes_and_increments(): response_cost=0.25, ) - counter = counter_cache.in_memory_cache.get_cache( - key=f"spend:key:{hashed_token}" - ) + counter = counter_cache.in_memory_cache.get_cache(key=f"spend:key:{hashed_token}") assert counter == 5.75 finally: ps.user_api_key_cache = original_key_cache @@ -7279,9 +6911,7 @@ async def test_increment_spend_counters_team_and_member(): team_counter = counter_cache.in_memory_cache.get_cache(key="spend:team:team-1") assert team_counter == 2.30 - member_counter = counter_cache.in_memory_cache.get_cache( - key="spend:team_member:user-1:team-1" - ) + member_counter = counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") assert member_counter == 1.30 finally: ps.user_api_key_cache = original_key_cache @@ -7339,14 +6969,10 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( increment=1.5, ) - fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": "team-9"} - ) + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. # Only the per-request delta (1.5) goes through INCRBYFLOAT. - fake_redis.async_set_cache.assert_awaited_once_with( - key="spend:team:team-9", value=42.0, nx=True - ) + fake_redis.async_set_cache.assert_awaited_once_with(key="spend:team:team-9", value=42.0, nx=True) writes = [(c["key"], c["value"]) for c in recorded_increments] assert writes == [("spend:team:team-9", 1.5)] finally: @@ -7415,9 +7041,7 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( return row fake_prisma = MagicMock() - fake_prisma.db.litellm_teamtable.find_unique = AsyncMock( - side_effect=slow_find_unique - ) + fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(side_effect=slow_find_unique) pod_a = DualCache() pod_a.redis_cache = fake_redis @@ -7450,11 +7074,7 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( # (winner) and one was rejected (loser). assert db_read_count == 2 assert fake_redis.async_set_cache.await_count == 2 - nx_writes = [ - call - for call in fake_redis.async_set_cache.await_args_list - if call.kwargs.get("nx") is True - ] + nx_writes = [call for call in fake_redis.async_set_cache.await_args_list if call.kwargs.get("nx") is True] assert len(nx_writes) == 2 assert sorted(set_results) == [ False, @@ -7463,9 +7083,7 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( # Loser path executed: after the winner's SET NX returned True, the # losing coalesced() call falls back to async_get_cache to read the # winner's value rather than re-seeding. - assert ( - get_after_set_count >= 1 - ), "loser branch (else: read back winner's value) was never exercised" + assert get_after_set_count >= 1, "loser branch (else: read back winner's value) was never exercised" @pytest.mark.asyncio @@ -7487,14 +7105,10 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) fake_prisma.db.litellm_endusertable.find_unique = AsyncMock() fake_prisma.db.litellm_tagtable.find_unique = AsyncMock() - fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock( - return_value=org_row - ) + fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock(return_value=org_row) assert await SpendCounterReseed.from_db(fake_prisma, "spend:user:alice") == 17.0 - fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with( - where={"user_id": "alice"} - ) + fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "alice"}) assert ( await SpendCounterReseed.from_db( @@ -7509,9 +7123,7 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited() assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0 - fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( - where={"organization_id": "acme"} - ) + fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(where={"organization_id": "acme"}) @pytest.mark.asyncio @@ -7525,14 +7137,8 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): fake_prisma.db.litellm_verificationtoken.find_unique = AsyncMock() fake_prisma.db.litellm_teamtable.find_unique = AsyncMock() - assert ( - await SpendCounterReseed.from_db(fake_prisma, "spend:key:sk-abc:window:1h") - is None - ) - assert ( - await SpendCounterReseed.from_db(fake_prisma, "spend:team:team-1:window:1d") - is None - ) + assert await SpendCounterReseed.from_db(fake_prisma, "spend:key:sk-abc:window:1h") is None + assert await SpendCounterReseed.from_db(fake_prisma, "spend:team:team-1:window:1d") is None fake_prisma.db.litellm_verificationtoken.find_unique.assert_not_awaited() fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited() @@ -7568,9 +7174,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): where={"api_key": "key-window", "startTime": {"gte": window_start}}, sum={"spend": True}, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-window:window:1h" - ) == pytest.approx(2.75) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-window:window:1h") == pytest.approx(2.75) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7625,14 +7229,10 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): increment=1.5, ) - fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": "team-stale-local"} - ) + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. assert redis_store[counter_key] == pytest.approx(43.5) - assert counter_cache.in_memory_cache.get_cache( - key=counter_key - ) == pytest.approx(43.5) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(43.5) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7695,9 +7295,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): sum={"spend": True}, ) assert redis_store[counter_key] == pytest.approx(2.75) - assert counter_cache.in_memory_cache.get_cache( - key=counter_key - ) == pytest.approx(2.75) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.75) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7733,9 +7331,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() fake_prisma = MagicMock() fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[ - {"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}} - ] + return_value=[{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}] ) import litellm.proxy.proxy_server as ps @@ -7758,9 +7354,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() nx=True, ) assert redis_store[counter_key] == pytest.approx(3.25) - assert counter_cache.in_memory_cache.get_cache( - key=counter_key - ) == pytest.approx(3.25) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(3.25) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7786,12 +7380,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): increment=0.5, ) - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-invalid-window:window:not-a-duration" - ) - is None - ) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None finally: ps.spend_counter_cache = orig_counter @@ -7873,9 +7462,9 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): assert incremented_counters == ["spend:team:team-finalize-after-increments"] assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-finalize-after-increments" - ) == pytest.approx(0.25) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-finalize-after-increments") == pytest.approx( + 0.25 + ) finally: ps.spend_counter_cache = orig_counter ps.user_api_key_cache = orig_user @@ -7919,9 +7508,7 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation(): ) assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-finalize-none-cost" - ) == pytest.approx(0.0) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-finalize-none-cost") == pytest.approx(0.0) finally: ps.spend_counter_cache = orig_counter @@ -7971,9 +7558,7 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( assert budget_reservation["finalized"] is True # counter reseeded to the authoritative DB value, not deleted/left None # and not double-counted via a direct increment - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-bad-reserved-counter" - ) == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.6) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8002,12 +7587,8 @@ async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure( increment=0.5, ) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None - ) - fake_redis.async_delete_cache.assert_awaited_once_with( - key="spend:team:redis-fail" - ) + assert counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None + fake_redis.async_delete_cache.assert_awaited_once_with(key="spend:team:redis-fail") finally: ps.spend_counter_cache = orig_counter @@ -8053,16 +7634,13 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): fallback_spend=30.0, ) assert spend == 362.0, ( - f"expected DB reseed to return 362.0, got {spend} " - f"(fallback would have returned 30.0 and caused bypass)" + f"expected DB reseed to return 362.0, got {spend} (fallback would have returned 30.0 and caused bypass)" ) # Counter warmed via SET NX so subsequent reads are fast. assert ("spend:team_member:user-1:team-1", 362.0, True) in [ (s["key"], s["value"], s["nx"]) for s in recorded_seeds ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:team_member:user-1:team-1" - ) == pytest.approx(362.0) + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == pytest.approx(362.0) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8147,9 +7725,7 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - side_effect=slow_find_unique - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=slow_find_unique) import litellm.proxy.proxy_server as ps @@ -8158,15 +7734,10 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): ps.prisma_client = fake_prisma try: results = await _asyncio.gather( - *[ - get_current_spend(counter_key=counter_key, fallback_spend=0.0) - for _ in range(5) - ] + *[get_current_spend(counter_key=counter_key, fallback_spend=0.0) for _ in range(5)] ) assert results == [100.0] * 5, f"all callers should see DB value, got {results}" - assert ( - db_call_count == 1 - ), f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}" + assert db_call_count == 1, f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}" finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8203,9 +7774,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): counter_key="spend:team_member:user-1:team-after-reset", fallback_spend=42.0, ) - assert ( - spend == 0.0 - ), f"DB authoritative 0 must override stale fallback 42, got {spend}" + assert spend == 0.0, f"DB authoritative 0 must override stale fallback 42, got {spend}" finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8263,9 +7832,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - side_effect=slow_find_unique - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=slow_find_unique) import litellm.proxy.proxy_server as ps @@ -8287,9 +7854,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): ), get_current_spend(counter_key=counter_key, fallback_spend=0.0), ) - assert ( - db_call_count == 1 - ), f"expected 1 DB query for concurrent read+write+read, got {db_call_count}" + assert db_call_count == 1, f"expected 1 DB query for concurrent read+write+read, got {db_call_count}" # Read-path callers see the warmed counter; the write path's # increment may or may not have landed by then, so accept either # the seeded value or seeded+increment. @@ -8325,9 +7890,7 @@ async def test_reseed_locks_dict_is_bounded(): try: for i in range(7): await SpendCounterReseed._get_lock(f"spend:key:test-key-{i}") - assert ( - len(SpendCounterReseed._locks) == 5 - ), f"got {len(SpendCounterReseed._locks)}" + assert len(SpendCounterReseed._locks) == 5, f"got {len(SpendCounterReseed._locks)}" # Oldest two evicted assert "spend:key:test-key-0" not in SpendCounterReseed._locks assert "spend:key:test-key-1" not in SpendCounterReseed._locks @@ -8384,9 +7947,7 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): return row fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - side_effect=find_unique - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=find_unique) import litellm.proxy.proxy_server as ps @@ -8399,9 +7960,7 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): # Second call: cache should be warmed at 0, no second DB query. spend2 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) assert spend1 == 0.0 and spend2 == 0.0 - assert ( - db_call_count == 1 - ), f"second read should hit warmed cache, got {db_call_count} DB queries" + assert db_call_count == 1, f"second read should hit warmed cache, got {db_call_count} DB queries" assert redis_store.get(counter_key) == 0.0, "cache must be warmed at 0" finally: ps.spend_counter_cache = orig_counter @@ -8464,9 +8023,7 @@ def _update_config_setup(monkeypatch): def _install(initial_rows=None, store_model_in_db=True): prisma = _FakePrismaClient(initial_rows=initial_rows) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) - monkeypatch.setattr( - "litellm.proxy.proxy_server.store_model_in_db", store_model_in_db - ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", store_model_in_db) monkeypatch.setattr( "litellm.proxy.proxy_server.encrypt_value_helper", lambda value, **_: f"enc:{value}", @@ -8477,9 +8034,7 @@ def _update_config_setup(monkeypatch): ) from litellm.proxy.proxy_server import proxy_config as real_proxy_config - monkeypatch.setattr( - real_proxy_config, "add_deployment", AsyncMock(return_value=None) - ) + monkeypatch.setattr(real_proxy_config, "add_deployment", AsyncMock(return_value=None)) original_overrides = app.dependency_overrides.copy() app.dependency_overrides[auth_dep] = lambda: UserAPIKeyAuth( @@ -8514,19 +8069,13 @@ def test_update_config_writes_only_sent_section(_update_config_setup): assert resp.status_code == 200 written = {name for name, _ in prisma.db.litellm_config.upsert_calls} assert written == {"general_settings"} - assert prisma.db.litellm_config.rows["litellm_settings"] == { - "drop_params": True - } - assert prisma.db.litellm_config.rows["environment_variables"] == { - "FOO": "enc:bar" - } + assert prisma.db.litellm_config.rows["litellm_settings"] == {"drop_params": True} + assert prisma.db.litellm_config.rows["environment_variables"] == {"FOO": "enc:bar"} finally: restore() -def test_update_config_env_var_round_trip_not_double_encrypted( - _update_config_setup, monkeypatch -): +def test_update_config_env_var_round_trip_not_double_encrypted(_update_config_setup, monkeypatch): """Endpoint-level regression for the /config/update double-encryption bug. The Admin UI reads config back via /get/config/callbacks (which returns @@ -8539,16 +8088,12 @@ def test_update_config_env_var_round_trip_not_double_encrypted( code this stored "enc:enc:..."; the assertions below would fail there. """ - def _fake_decrypt( - value, key=None, exception_type="error", return_original_value=False - ): + def _fake_decrypt(value, key=None, exception_type="error", return_original_value=False): if isinstance(value, str) and value.startswith("enc:"): return value[len("enc:") :] return value if return_original_value else None - monkeypatch.setattr( - "litellm.proxy.proxy_server.decrypt_value_helper", _fake_decrypt - ) + monkeypatch.setattr("litellm.proxy.proxy_server.decrypt_value_helper", _fake_decrypt) client, prisma, restore = _update_config_setup( initial_rows={"environment_variables": {"PREEXISTING_KEY": "enc:keepme"}} @@ -8566,21 +8111,14 @@ def test_update_config_env_var_round_trip_not_double_encrypted( # UI round-trip: re-POST the stored ciphertext (no field change). resp = client.post( "/config/update", - json={ - "environment_variables": { - "LANGFUSE_SECRET_KEY": stored["LANGFUSE_SECRET_KEY"] - } - }, + json={"environment_variables": {"LANGFUSE_SECRET_KEY": stored["LANGFUSE_SECRET_KEY"]}}, ) assert resp.status_code == 200 stored = prisma.db.litellm_config.rows["environment_variables"] # The bug: this would be "enc:enc:sk-secret". The fix keeps it single. assert stored["LANGFUSE_SECRET_KEY"] == "enc:sk-secret" - assert ( - _fake_decrypt(stored["LANGFUSE_SECRET_KEY"], return_original_value=True) - == "sk-secret" - ) + assert _fake_decrypt(stored["LANGFUSE_SECRET_KEY"], return_original_value=True) == "sk-secret" # Untouched key preserved byte-for-byte (only sent keys rewritten). assert stored["PREEXISTING_KEY"] == "enc:keepme" @@ -8595,14 +8133,9 @@ def test_update_config_can_flip_store_model_in_db_when_currently_false( False, blocking the very request that would flip it to True.""" client, prisma, restore = _update_config_setup(store_model_in_db=False) try: - resp = client.post( - "/config/update", json={"general_settings": {"store_model_in_db": True}} - ) + resp = client.post("/config/update", json={"general_settings": {"store_model_in_db": True}}) assert resp.status_code == 200 - assert ( - prisma.db.litellm_config.rows["general_settings"]["store_model_in_db"] - is True - ) + assert prisma.db.litellm_config.rows["general_settings"]["store_model_in_db"] is True finally: restore() @@ -8635,9 +8168,7 @@ def test_update_config_litellm_settings_request_wins_for_non_callback_keys( } ) try: - resp = client.post( - "/config/update", json={"litellm_settings": {"drop_params": False}} - ) + resp = client.post("/config/update", json={"litellm_settings": {"drop_params": False}}) assert resp.status_code == 200 stored = prisma.db.litellm_config.rows["litellm_settings"] assert stored["drop_params"] is False @@ -8733,9 +8264,7 @@ class TestLazyFeaturesNotImportedAtStartup: from litellm.proxy._lazy_features import LAZY_FEATURES - proxy_server_src = ( - Path(__file__).resolve().parents[3] / "litellm/proxy/proxy_server.py" - ).read_text() + proxy_server_src = (Path(__file__).resolve().parents[3] / "litellm/proxy/proxy_server.py").read_text() leaks = [] for feat in LAZY_FEATURES: @@ -8840,9 +8369,7 @@ class TestLazyFeatureMiddleware: ("/api/v1", "/api/v1/unrelated", False, "unrelated path under root"), ], ) - async def test_root_path_handling( - self, monkeypatch, server_root_path, request_path, should_load, case - ): + async def test_root_path_handling(self, monkeypatch, server_root_path, request_path, should_load, case): """ The middleware must strip SERVER_ROOT_PATH before prefix-matching so lazy features load under deployments that set a server root path, @@ -8952,9 +8479,7 @@ class TestLazyFeatureMiddleware: ) await asyncio.gather(hit(), hit(), hit(), hit(), hit()) - assert loads == [ - "json" - ], f"expected one registration despite concurrent first hits, got {loads}" + assert loads == ["json"], f"expected one registration despite concurrent first hits, got {loads}" @pytest.mark.asyncio async def test_failing_import_does_not_loop(self): @@ -9004,9 +8529,9 @@ class TestLazyFeatureMiddleware: receive, send, ) - assert attempts == [ - "called" - ], f"failing register_fn should be invoked once, not on every request; got {attempts}" + assert attempts == ["called"], ( + f"failing register_fn should be invoked once, not on every request; got {attempts}" + ) @pytest.mark.asyncio @@ -9074,9 +8599,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - return_value=MagicMock(spend=999.0) - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=MagicMock(spend=999.0)) import litellm.proxy.proxy_server as ps @@ -9086,8 +8609,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(): try: spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) assert spend == 42.0, ( - f"expected in-memory fallback 42.0 on Redis error, got {spend} " - f"(should not have hit DB when Redis errored)" + f"expected in-memory fallback 42.0 on Redis error, got {spend} (should not have hit DB when Redis errored)" ) # DB query should NOT have fired - in-memory short-circuits. fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() @@ -9110,9 +8632,7 @@ def test_realtime_websocket_route_aliases_registered(): from litellm.proxy.proxy_server import app from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes - websocket_paths = { - route.path for route in app.routes if isinstance(route, WebSocketRoute) - } + websocket_paths = {route.path for route in app.routes if isinstance(route, WebSocketRoute)} openai_routes = LiteLLMRoutes.openai_routes.value for expected in ("/openai/v1/realtime", "/v1/realtime", "/realtime"): @@ -9174,8 +8694,7 @@ class TestTransformRequestBannedParams: }, ) assert response.status_code == 400, ( - f"Expected 400 for banned param '{banned}', " - f"got {response.status_code}: {response.json()}" + f"Expected 400 for banned param '{banned}', got {response.status_code}: {response.json()}" ) @@ -9201,13 +8720,8 @@ class TestSortModelsByDisplayName: {"model_name": "gpt-4o", "model_info": {}}, ] - sorted_models = _sort_models( - all_models=models, sort_by="model_name", sort_order="asc" - ) - displayed_order = [ - m["model_info"].get("team_public_model_name") or m["model_name"] - for m in sorted_models - ] + sorted_models = _sort_models(all_models=models, sort_by="model_name", sort_order="asc") + displayed_order = [m["model_info"].get("team_public_model_name") or m["model_name"] for m in sorted_models] assert displayed_order == [ "anthropic/claude", "claude-haiku-4-5", @@ -9226,13 +8740,8 @@ class TestSortModelsByDisplayName: {"model_name": "gpt-4o", "model_info": {}}, ] - sorted_models = _sort_models( - all_models=models, sort_by="model_name", sort_order="desc" - ) - displayed_order = [ - m["model_info"].get("team_public_model_name") or m["model_name"] - for m in sorted_models - ] + sorted_models = _sort_models(all_models=models, sort_by="model_name", sort_order="desc") + displayed_order = [m["model_info"].get("team_public_model_name") or m["model_name"] for m in sorted_models] assert displayed_order == [ "zeta/model", "gpt-4o", @@ -9250,9 +8759,7 @@ class TestSortModelsByDisplayName: {"model_name": "beta", "model_info": {}}, ] - sorted_models = _sort_models( - all_models=models, sort_by="model_name", sort_order="asc" - ) + sorted_models = _sort_models(all_models=models, sort_by="model_name", sort_order="asc") assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"] @@ -9274,9 +8781,7 @@ class TestDeleteDeploymentSync: mock_router.delete_deployment.return_value = MagicMock() with patch("litellm.proxy.proxy_server.llm_router", mock_router): - with patch.object( - proxy_config, "get_config", AsyncMock(return_value={"model_list": []}) - ): + with patch.object(proxy_config, "get_config", AsyncMock(return_value={"model_list": []})): count = await proxy_config._delete_deployment(db_models=[]) mock_router.delete_deployment.assert_called_once_with(id="model-id-to-evict") @@ -9297,9 +8802,7 @@ class TestDeleteDeploymentSync: with patch("litellm.proxy.proxy_server.llm_router", mock_router): with patch.object(proxy_config, "get_config", AsyncMock(return_value={})): - await proxy_config._update_llm_router( - new_models=None, proxy_logging_obj=MagicMock() - ) + await proxy_config._update_llm_router(new_models=None, proxy_logging_obj=MagicMock()) mock_router.delete_deployment.assert_not_called() mock_router.upsert_deployment.assert_not_called() @@ -9316,15 +8819,11 @@ class TestDeleteDeploymentSync: proxy_config = ProxyConfig() mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( - side_effect=Exception("DB connection lost") - ) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=Exception("DB connection lost")) result = await proxy_config._get_models_from_db(prisma_client=mock_prisma) - assert ( - result is None - ), f"Expected None on DB failure to signal fetch error, got {result!r}" + assert result is None, f"Expected None on DB failure to signal fetch error, got {result!r}" def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): @@ -9606,9 +9105,18 @@ def test_general_settings_ui_defaults_unchanged_for_existing_fields(): _general_settings_ui_litellm_default, ) - assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["budget_exceeded_throttle_percentage"]) is None - assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["enable_anthropic_prompt_caching"]) is False - assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["anthropic_prompt_caching_ttl"]) is None + assert ( + _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["budget_exceeded_throttle_percentage"]) + is None + ) + assert ( + _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["enable_anthropic_prompt_caching"]) + is False + ) + assert ( + _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["anthropic_prompt_caching_ttl"]) + is None + ) @pytest.mark.parametrize( @@ -9874,16 +9382,10 @@ def test_preserve_redacted_plugin_keys_keeps_stored_credential(): existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] - redacted = _preserve_redacted_plugin_keys( - [{"name": "p1", "url": "https://p1-new", "plugin_key": "***"}], existing - ) - assert redacted == [ - {"name": "p1", "url": "https://p1-new", "plugin_key": "sk-real-1"} - ] + redacted = _preserve_redacted_plugin_keys([{"name": "p1", "url": "https://p1-new", "plugin_key": "***"}], existing) + assert redacted == [{"name": "p1", "url": "https://p1-new", "plugin_key": "sk-real-1"}] - blanked = _preserve_redacted_plugin_keys( - [{"name": "p1", "url": "https://p1", "plugin_key": ""}], existing - ) + blanked = _preserve_redacted_plugin_keys([{"name": "p1", "url": "https://p1", "plugin_key": ""}], existing) assert blanked[0]["plugin_key"] == "sk-real-1" @@ -9893,14 +9395,10 @@ def test_preserve_redacted_plugin_keys_sets_new_and_drops_orphan_placeholder(): existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] - rotated = _preserve_redacted_plugin_keys( - [{"name": "p1", "url": "https://p1", "plugin_key": "sk-new"}], existing - ) + rotated = _preserve_redacted_plugin_keys([{"name": "p1", "url": "https://p1", "plugin_key": "sk-new"}], existing) assert rotated[0]["plugin_key"] == "sk-new" - new_plugin = _preserve_redacted_plugin_keys( - [{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing - ) + new_plugin = _preserve_redacted_plugin_keys([{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing) assert "plugin_key" not in new_plugin[0] @@ -9933,9 +9431,7 @@ def _config_field_info_client(monkeypatch, user_role): mock_prisma = MagicMock() mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="u", user_role=user_role - ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=user_role) return TestClient(app) @@ -9946,9 +9442,7 @@ def test_config_field_info_redacts_secrets_for_view_only_admin(monkeypatch): is not a FULL PROXY_ADMIN, while non-secret fields stay readable.""" from litellm.proxy._types import LitellmUserRoles - client = _config_field_info_client( - monkeypatch, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ) + client = _config_field_info_client(monkeypatch, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) try: for secret_field in ("master_key", "database_url", "pass_through_endpoints"): resp = client.get("/config/field/info", params={"field_name": secret_field}) @@ -9958,9 +9452,7 @@ def test_config_field_info_redacts_secrets_for_view_only_admin(monkeypatch): assert "secret" not in str(body["field_value"]) assert "p4ssw0rd" not in str(body["field_value"]) - resp = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + resp = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert resp.status_code == 200, resp.text assert resp.json()["field_value"] == 100 finally: @@ -9978,14 +9470,9 @@ def test_config_field_info_returns_raw_secrets_for_full_admin(monkeypatch): assert resp.status_code == 200, resp.text assert resp.json()["field_value"] == "sk-super-secret-master" - resp = client.get( - "/config/field/info", params={"field_name": "pass_through_endpoints"} - ) + resp = client.get("/config/field/info", params={"field_name": "pass_through_endpoints"}) assert resp.status_code == 200, resp.text - assert ( - resp.json()["field_value"][0]["headers"]["Authorization"] - == "Bearer sk-upstream-secret" - ) + assert resp.json()["field_value"][0]["headers"]["Authorization"] == "Bearer sk-upstream-secret" finally: app.dependency_overrides.clear() @@ -10227,9 +9714,7 @@ async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatc user_role=LitellmUserRoles.PROXY_ADMIN, ) await delete_config_general_settings( - data=ConfigFieldDelete( - field_name="max_parallel_requests", config_type="general_settings" - ), + data=ConfigFieldDelete(field_name="max_parallel_requests", config_type="general_settings"), user_api_key_dict=admin, ) # Audit is scheduled via asyncio.create_task; yield so it runs. @@ -10252,9 +9737,7 @@ def test_update_config_audits_every_written_section(_update_config_setup, monkey is the row that holds default_internal_user_params ("default user settings").""" import litellm.proxy.proxy_server as proxy_server_module - client, prisma, restore = _update_config_setup( - initial_rows={"litellm_settings": {"drop_params": True}} - ) + client, prisma, restore = _update_config_setup(initial_rows={"litellm_settings": {"drop_params": True}}) audit_create = AsyncMock() prisma.db.litellm_auditlog.create = audit_create monkeypatch.setattr(proxy_server_module, "premium_user", True) @@ -10265,17 +9748,14 @@ def test_update_config_audits_every_written_section(_update_config_setup, monkey json={ "general_settings": {"store_prompts_in_spend_logs": True}, "environment_variables": {"FOO": "bar"}, - "litellm_settings": { - "default_internal_user_params": {"max_budget": 10} - }, + "litellm_settings": {"default_internal_user_params": {"max_budget": 10}}, "router_settings": {"routing_strategy": "latency-based-routing"}, }, ) assert resp.status_code == 200, resp.text audited = { - call.kwargs["data"]["object_id"]: call.kwargs["data"]["action"] - for call in audit_create.await_args_list + call.kwargs["data"]["object_id"]: call.kwargs["data"]["action"] for call in audit_create.await_args_list } assert audited == { "general_settings": "updated", @@ -10287,20 +9767,14 @@ def test_update_config_audits_every_written_section(_update_config_setup, monkey assert call.kwargs["data"]["table_name"] == "LiteLLM_Config" assert call.kwargs["data"]["changed_by"] == "test_admin" - ls_call = next( - c - for c in audit_create.await_args_list - if c.kwargs["data"]["object_id"] == "litellm_settings" - ) + ls_call = next(c for c in audit_create.await_args_list if c.kwargs["data"]["object_id"] == "litellm_settings") after = json.loads(ls_call.kwargs["data"]["updated_values"]) assert after["default_internal_user_params"] == {"max_budget": 10} finally: restore() -def test_delete_callback_audits_litellm_settings_deletion( - _update_config_setup, monkeypatch -): +def test_delete_callback_audits_litellm_settings_deletion(_update_config_setup, monkeypatch): """/config/callback/delete must emit a deleted audit row for litellm_settings capturing the success_callback list before and after removal.""" import litellm.proxy.proxy_server as proxy_server_module @@ -10316,19 +9790,11 @@ def test_delete_callback_audits_litellm_settings_deletion( monkeypatch.setattr( real_proxy_config, "get_config", - AsyncMock( - return_value={ - "litellm_settings": {"success_callback": ["langfuse", "datadog"]} - } - ), - ) - monkeypatch.setattr( - real_proxy_config, "save_config", AsyncMock(return_value=None) + AsyncMock(return_value={"litellm_settings": {"success_callback": ["langfuse", "datadog"]}}), ) + monkeypatch.setattr(real_proxy_config, "save_config", AsyncMock(return_value=None)) try: - resp = client.post( - "/config/callback/delete", json={"callback_name": "datadog"} - ) + resp = client.post("/config/callback/delete", json={"callback_name": "datadog"}) assert resp.status_code == 200, resp.text audit_create.assert_awaited_once() @@ -10357,24 +9823,16 @@ def test_delete_callback_audits_before_reload_failure(_update_config_setup, monk monkeypatch.setattr( real_proxy_config, "get_config", - AsyncMock( - return_value={ - "litellm_settings": {"success_callback": ["langfuse", "datadog"]} - } - ), - ) - monkeypatch.setattr( - real_proxy_config, "save_config", AsyncMock(return_value=None) + AsyncMock(return_value={"litellm_settings": {"success_callback": ["langfuse", "datadog"]}}), ) + monkeypatch.setattr(real_proxy_config, "save_config", AsyncMock(return_value=None)) monkeypatch.setattr( real_proxy_config, "add_deployment", AsyncMock(side_effect=RuntimeError("reload failed")), ) try: - resp = client.post( - "/config/callback/delete", json={"callback_name": "datadog"} - ) + resp = client.post("/config/callback/delete", json={"callback_name": "datadog"}) assert resp.status_code == 500, resp.text audit_create.assert_awaited_once() @@ -10385,9 +9843,7 @@ def test_delete_callback_audits_before_reload_failure(_update_config_setup, monk restore() -def test_update_config_redacts_all_environment_variable_values( - _update_config_setup, monkeypatch -): +def test_update_config_redacts_all_environment_variable_values(_update_config_setup, monkeypatch): """environment_variables hold credentials under arbitrary uppercase keys (DATABASE_URL) that key-name secret matching misses, so every value in the section must be redacted before the audit row is written; a plaintext @@ -10397,11 +9853,7 @@ def test_update_config_redacts_all_environment_variable_values( # DATABASE_URL is the bug class: an uppercase env key that key-name secret # matching does NOT flag, so only whole-section value redaction protects it. client, prisma, restore = _update_config_setup( - initial_rows={ - "environment_variables": { - "DATABASE_URL": "enc:postgresql://OLDsecret@old.host:5432/db" - } - } + initial_rows={"environment_variables": {"DATABASE_URL": "enc:postgresql://OLDsecret@old.host:5432/db"}} ) audit_create = AsyncMock() prisma.db.litellm_auditlog.create = audit_create @@ -10420,9 +9872,7 @@ def test_update_config_redacts_all_environment_variable_values( assert resp.status_code == 200, resp.text env_call = next( - c - for c in audit_create.await_args_list - if c.kwargs["data"]["object_id"] == "environment_variables" + c for c in audit_create.await_args_list if c.kwargs["data"]["object_id"] == "environment_variables" ) data = env_call.kwargs["data"] @@ -10586,11 +10036,7 @@ def test_init_coordination_redis_startup_nodes_builds_cluster_client(): """A coordination_redis block with startup_nodes must construct a cluster client, so cluster-aware consumers (v3 rate limiter) take the cluster path.""" usage_cache, _, _ = _run_init_coordination_redis( - config={ - "general_settings": { - "coordination_redis": {"startup_nodes": [{"host": "node-1", "port": 7000}]} - } - }, + config={"general_settings": {"coordination_redis": {"startup_nodes": [{"host": "node-1", "port": 7000}]}}}, ) assert isinstance(usage_cache, _EnvBuiltClusterCache)