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)