From b079bf5a51ddad9c7d5f89d515d3751ea230e5bf Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 12 Nov 2026 04:40:00 +0900 Subject: [PATCH] refactor(proxy): let the gunicorn server take a private listener (#50871) Extract serving_port() and configure_database() as they were inline, turn the LITELLM_RUST parser into a reusable RustSwitch type, and reformat tests/unit/proxy/test_proxy_cli.py to the 120 column limit. run_gunicorn_server takes private_listener, an inherited socket or an assigned loopback port. With one, gunicorn binds it instead of host:port and trusts no peer's forwarding headers, since every peer on it is the local process in front. Both are gunicorn's own settings, bind and forwarded_allow_ips, so the uvicorn worker is unchanged. The reload options take factory. No caller passes either yet, so every entry point serves as before. --- gateway/launch.py | 6 +- litellm/proxy/dev_reload.py | 7 +- litellm/proxy/proxy_cli.py | 16 +- litellm/proxy/server_listener.py | 38 ++ litellm/rust_bridge/configuration.py | 37 +- tests/unit/proxy/test_dev_reload.py | 41 +- tests/unit/proxy/test_proxy_cli.py | 667 +++++++------------ tests/unit/proxy/test_server_listener.py | 28 + tests/unit/rust_bridge/test_configuration.py | 13 +- 9 files changed, 394 insertions(+), 459 deletions(-) create mode 100644 litellm/proxy/server_listener.py create mode 100644 tests/unit/proxy/test_server_listener.py diff --git a/gateway/launch.py b/gateway/launch.py index 59f7121e513..b9303fd9447 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -75,7 +75,7 @@ def _serve(argv: Sequence[str]) -> None: uvicorn_main(tuple(argv), prog_name="uvicorn") -def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: +def configure_database() -> None: from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url @@ -86,6 +86,10 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}") if pooled_url is not None: export_pooled_database_url(pooled_url) + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + configure_database() serve(uvicorn_argv(argv, os.environ)) diff --git a/litellm/proxy/dev_reload.py b/litellm/proxy/dev_reload.py index 500f581b34d..7b057271cb6 100644 --- a/litellm/proxy/dev_reload.py +++ b/litellm/proxy/dev_reload.py @@ -13,7 +13,7 @@ from socket import socket from typing import Final, Literal from pydantic import ConfigDict, ImportString, JsonValue, TypeAdapter, with_config -from typing_extensions import ReadOnly, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from uvicorn import Config, Server from uvicorn import config as uvicorn_config from uvicorn.supervisors import ChangeReload, Multiprocess @@ -42,6 +42,7 @@ class ServerOptions(TypedDict, total=False): @with_config(ConfigDict(extra="allow")) class ReloadOptions(ServerOptions): app: Required[ReadOnly[str]] + factory: NotRequired[ReadOnly[bool]] reload: ReadOnly[bool] reload_dirs: ReadOnly[tuple[str, ...]] reload_includes: ReadOnly[tuple[str, ...]] @@ -92,7 +93,9 @@ def load_worker_app(import_path: str, worker_app_factory: bool) -> Callable[..., def pool_config(options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None) -> Config: overrides: Final = worker_overrides or {} worker_app: Final = TypeAdapter(str).validate_python(overrides.get("app", options["app"])) - worker_app_factory: Final = TypeAdapter(bool).validate_python(overrides.get("factory", False)) + worker_app_factory: Final = TypeAdapter(bool).validate_python( + overrides.get("factory", options.get("factory", False)) + ) return Config( partial(load_worker_app, worker_app, worker_app_factory), factory=True, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 41aedf4b5cd..b238cefd24f 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -31,6 +31,7 @@ from litellm.proxy.db.pgbouncer import ( start_in_container_pgbouncer, ) from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper +from litellm.proxy.server_listener import PrivateListener, gunicorn_overrides from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: @@ -686,6 +687,7 @@ class ProxyInitializationHelpers: max_requests_before_restart: int | None = None, max_requests_before_restart_jitter: int | None = None, max_request_header_size: int | None = None, + private_listener: PrivateListener | None = None, ): """ Run litellm with `gunicorn` @@ -766,6 +768,7 @@ class ProxyInitializationHelpers: "accesslog": "-", # Log to stdout "timeout": 600, # default to very high number, bedrock/anthropic.claude-v2:1 can take 30+ seconds for the 1st chunk to come in "access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s', + **gunicorn_overrides(private_listener), } # Optional: recycle workers after N requests to mitigate memory growth @@ -821,7 +824,7 @@ class ProxyInitializationHelpers: _run_ollama_serve = run_ollama_serve @staticmethod - def is_port_in_use(port): + def is_port_in_use(port: int): import socket with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: @@ -829,6 +832,10 @@ class ProxyInitializationHelpers: _is_port_in_use = is_port_in_use + @staticmethod + def serving_port(port: int) -> int: + return random.randint(1024, 49152) if port == 4000 and ProxyInitializationHelpers.is_port_in_use(port) else port + @staticmethod def get_loop_type(): """Helper function to determine the event loop type based on platform""" @@ -1228,8 +1235,8 @@ class ProxyInitializationHelpers: ) def run_server( cli_args, - host, - port, + host: str, + port: int, api_base, api_version, model, @@ -1717,8 +1724,7 @@ def run_server( ) sys.exit(1) export_pooled_database_url(pooled_database_url) - if port == 4000 and ProxyInitializationHelpers.is_port_in_use(port): - port = random.randint(1024, 49152) + port = ProxyInitializationHelpers.serving_port(port) if prometheus_metrics_port == port: raise click.UsageError("--prometheus_metrics_port must differ from --port") diff --git a/litellm/proxy/server_listener.py b/litellm/proxy/server_listener.py new file mode 100644 index 00000000000..bf57fce4ebd --- /dev/null +++ b/litellm/proxy/server_listener.py @@ -0,0 +1,38 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, TypeAlias + +from typing_extensions import assert_never + +LOOPBACK_HOST: Final = "127.0.0.1" +NO_TRUSTED_PROXIES: Final = "" + + +@dataclass(frozen=True, slots=True) +class InheritedFd: + fd: int + + +@dataclass(frozen=True, slots=True) +class AssignedPort: + port: int + + +PrivateListener: TypeAlias = InheritedFd | AssignedPort + + +def _gunicorn_bind(listener: PrivateListener) -> str: + match listener: + case InheritedFd(fd): + return f"fd://{fd}" + case AssignedPort(port): + return f"{LOOPBACK_HOST}:{port}" + case _: + assert_never(listener) + + +def gunicorn_overrides(listener: PrivateListener | None) -> Mapping[str, str]: + if listener is None: + return MappingProxyType({}) + return MappingProxyType({"bind": _gunicorn_bind(listener), "forwarded_allow_ips": NO_TRUSTED_PROXIES}) diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 152ba632996..dedbf857a07 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -3,13 +3,29 @@ from __future__ import annotations import os from enum import Enum, auto from functools import lru_cache -from typing import Final +from typing import Annotated, Final, TypeAlias -from pydantic import TypeAdapter, ValidationError +from pydantic import BeforeValidator, TypeAdapter, ValidationError from typing_extensions import assert_never -_GLOBAL_ENV_NAME: Final = "LITELLM_RUST" -_ENV_BOOL: Final = TypeAdapter(bool) +ENV_NAME: Final = "LITELLM_RUST" + + +def _blank_as_unset(value: object) -> object: + stripped: Final = value.strip() if isinstance(value, str) else value + return None if stripped == "" else stripped + + +RustSwitch: TypeAlias = Annotated[bool | None, BeforeValidator(_blank_as_unset)] +_RUST_SWITCH: Final = TypeAdapter[bool | None](RustSwitch) + + +@lru_cache(maxsize=16) +def _switch_or_unset(raw: str | None) -> bool | None: + try: + return _RUST_SWITCH.validate_python(raw) + except ValidationError: + return None class Rollout(Enum): @@ -33,17 +49,6 @@ class _RustConfiguration: _CONFIGURATION: Final = _RustConfiguration() -@lru_cache(maxsize=16) -def _parse_env_bool(value: str | None) -> bool | None: - """`LITELLM_RUST` as a bool; cached by raw value because `decision` runs per tokenizer call.""" - if value is None: - return None - try: - return _ENV_BOOL.validate_python(value.strip()) - except ValidationError: - return None - - def decide( rollout: Rollout, *, @@ -72,7 +77,7 @@ def decision(rollout: Rollout) -> Decision: return decide( rollout, process_override=_CONFIGURATION.override, - environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)), + environment_override=_switch_or_unset(os.getenv(ENV_NAME)), ) diff --git a/tests/unit/proxy/test_dev_reload.py b/tests/unit/proxy/test_dev_reload.py index 60a7c53820b..6650c1cd3e3 100644 --- a/tests/unit/proxy/test_dev_reload.py +++ b/tests/unit/proxy/test_dev_reload.py @@ -86,6 +86,39 @@ def test_a_broken_factory_exits_the_worker_as_a_startup_failure( assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE +@pytest.mark.parametrize("worker_overrides", [None, {"factory": True}]) +def test_a_factory_flag_in_the_options_reaches_the_workers( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, worker_overrides: dict[str, object] | None +) -> None: + (tmp_path / "reload_factory_flag_app.py").write_text( + "async def served(scope, receive, send):\n" + " await send({'type': 'http.response.start', 'status': 204, 'headers': []})\n" + " await send({'type': 'http.response.body', 'body': b'', 'more_body': False})\n" + "\n" + "def factory():\n" + " return served\n" + ) + monkeypatch.syspath_prepend(str(tmp_path)) + config: Final = pool_config( + RELOAD_OPTIONS.validate_python({**_options("reload_factory_flag_app:factory"), "factory": True}), + workers=POOL_WORKERS, + worker_overrides=worker_overrides, + ) + config.load() + loaded_app: Final = config.loaded_app + assert loaded_app is not None + sent: Final[list[Message]] = [] + + async def receive() -> Message: + return {"type": "http.request", "body": b"", "more_body": False} + + async def send(message: Message) -> None: + sent.append(message) + + asyncio.run(loaded_app({"type": "http", "method": "GET", "path": "/", "headers": []}, receive, send)) + assert sent[0] == {"type": "http.response.start", "status": 204, "headers": []} + + @REQUIRES_UVICORN_STARTUP_FAILURE @pytest.mark.parametrize( ("module", "source"), @@ -118,7 +151,9 @@ def test_the_reload_pool_defaults_the_worker_healthcheck_and_keeps_an_explicit_o explicit: Final = RELOAD_WORKER_HEALTHCHECK_SECONDS + 7 assert pool_config(_options("app:app"), workers=2).timeout_worker_healthcheck == RELOAD_WORKER_HEALTHCHECK_SECONDS assert ( - pool_config({**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2).timeout_worker_healthcheck + pool_config( + {**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2 + ).timeout_worker_healthcheck == explicit ) @@ -142,7 +177,9 @@ def test_the_reload_pool_accepts_and_keeps_every_uvicorn_arg_the_proxy_cli_build def test_the_reload_pool_passes_a_uvicorn_arg_it_does_not_list_through_to_the_workers() -> None: - config: Final = pool_config(RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS) + config: Final = pool_config( + RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS + ) assert (config.backlog, config.reload) == (123, False) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 9a091959631..eec682f365f 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -34,6 +34,7 @@ from litellm._logging import turn_on_json from litellm.constants import DEFAULT_MAX_REQUEST_HEADER_SIZE, GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server +from litellm.proxy.server_listener import AssignedPort, InheritedFd from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT @@ -201,42 +202,30 @@ class TestProxyInitializationHelpers: # Execute with pytest.raises(ValueError, match="Invalid test value"): - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", True - ) + ProxyInitializationHelpers._run_test_chat_completion("localhost", 8000, "gpt-3.5-turbo", True) # Test with valid string test value - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", "http://test-url" - ) + ProxyInitializationHelpers._run_test_chat_completion("localhost", 8000, "gpt-3.5-turbo", "http://test-url") # Assert - mock_openai.assert_called_once_with( - api_key="My API Key", base_url="http://test-url" - ) + mock_openai.assert_called_once_with(api_key="My API Key", base_url="http://test-url") mock_client.chat.completions.create.assert_called() def test_get_default_unvicorn_init_args(self): # Test without log_config - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) assert args["app"] == "litellm.proxy.proxy_server:app" assert args["host"] == "localhost" assert args["port"] == 8000 assert args["server_header"] is False # Test with log_config - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, "log_config.json" - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, "log_config.json") assert args["log_config"] == "log_config.json" # Test with json_logs=True with patch("litellm.json_logs", True): - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) # When json_logs is True, log_config should be set to the JSON log config dict assert args["log_config"] is not None assert isinstance(args["log_config"], dict) @@ -244,15 +233,11 @@ class TestProxyInitializationHelpers: assert "formatters" in args["log_config"] # Test with keepalive_timeout - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, None, 60 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, None, 60) assert args["timeout_keep_alive"] == 60 # Test with both log_config and keepalive_timeout - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, "log_config.json", 120 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, "log_config.json", 120) assert args["log_config"] == "log_config.json" assert args["timeout_keep_alive"] == 120 @@ -383,14 +368,10 @@ class TestProxyInitializationHelpers: py_file = tmp_path / "module.py" py_file.write_text("x = 1\n") - applied = ProxyInitializationHelpers._patch_statreload_extra_paths( - [str(config_file)] - ) + applied = ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)]) assert applied is True - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)} assert config_file.resolve() in yielded_paths @@ -407,14 +388,10 @@ class TestProxyInitializationHelpers: env_file = tmp_path / ".env" env_file.write_text("FOO=bar\n") - applied = ProxyInitializationHelpers._patch_statreload_extra_paths( - [str(env_file)] - ) + applied = ProxyInitializationHelpers._patch_statreload_extra_paths([str(env_file)]) assert applied is True - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)} assert env_file.resolve() in yielded_paths @@ -426,10 +403,7 @@ class TestProxyInitializationHelpers: StatReload._litellm_patched_config_paths.clear() assert ProxyInitializationHelpers._patch_statreload_extra_paths([]) is False - assert ( - ProxyInitializationHelpers._patch_statreload_extra_paths([None, ""]) - is False - ) + assert ProxyInitializationHelpers._patch_statreload_extra_paths([None, ""]) is False def test_patch_statreload_extra_paths_is_idempotent(self, tmp_path): from pathlib import Path @@ -447,18 +421,14 @@ class TestProxyInitializationHelpers: for _ in range(3): ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)]) - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded = list(StatReload.iter_py_files(fake_self)) assert len(yielded) == len(set(map(str, yielded))) yielded_paths = {Path(p).resolve() for p in yielded} assert config_file.resolve() in yielded_paths assert py_file.resolve() in yielded_paths - def test_configure_dev_reload_watches_env_and_sets_override_flag( - self, tmp_path, monkeypatch - ): + def test_configure_dev_reload_watches_env_and_sets_override_flag(self, tmp_path, monkeypatch): from pathlib import Path from uvicorn.supervisors.statreload import StatReload @@ -475,9 +445,7 @@ class TestProxyInitializationHelpers: uvicorn_args: dict = {} with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning: - ProxyInitializationHelpers._configure_dev_reload( - uvicorn_args, str(config_file) - ) + ProxyInitializationHelpers._configure_dev_reload(uvicorn_args, str(config_file)) assert os.environ["LITELLM_DEV_ENV_HOT_RELOAD"] == "True" assert uvicorn_args["reload"] is True @@ -488,9 +456,7 @@ class TestProxyInitializationHelpers: assert "override" in warning_text assert ".env" in warning_text - fake_self = types.SimpleNamespace( - config=types.SimpleNamespace(reload_dirs=[tmp_path]) - ) + fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path])) yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)} assert env_file.resolve() in yielded_paths assert config_file.resolve() in yielded_paths @@ -514,17 +480,13 @@ class TestProxyInitializationHelpers: mock_app = MagicMock() # Execute - ProxyInitializationHelpers._init_hypercorn_server( - mock_app, "localhost", 8000, None, None, None - ) + ProxyInitializationHelpers._init_hypercorn_server(mock_app, "localhost", 8000, None, None, None) # Assert mock_asyncio_run.assert_called_once() # Test with SSL - ProxyInitializationHelpers._init_hypercorn_server( - mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE" - ) + ProxyInitializationHelpers._init_hypercorn_server(mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE") @patch("granian.Granian") @patch("builtins.print") @@ -693,10 +655,14 @@ class TestProxyInitializationHelpers: database_name_enc = urllib.parse.quote_plus(database_name) # Construct DATABASE_URL from the provided variables - database_url = f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + database_url = ( + f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + ) # Assert the correct URL was constructed with properly escaped characters - expected_url = "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + expected_url = ( + "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + ) assert database_url == expected_url # Test appending query parameters @@ -714,9 +680,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") # critical @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_skip_server_startup( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run, tmp_path: Path ): @@ -734,11 +698,7 @@ class TestProxyInitializationHelpers: ) # Remove DATABASE_URL/DIRECT_URL so the CLI doesn't attempt # real prisma operations when these are set in CI. - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) live_proxy_samples = tmp_path / "counter_123.db" live_proxy_samples.write_bytes(b"samples of a proxy that is still running") @@ -758,9 +718,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -771,18 +729,12 @@ class TestProxyInitializationHelpers: # --- skip startup --- result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "Skipping server startup" in result.output mock_uvicorn_run.assert_not_called() - result = runner.invoke( - run_server, ["--local", "--skip_server_startup", "--telemetry", "False"] - ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + result = runner.invoke(run_server, ["--local", "--skip_server_startup", "--telemetry", "False"]) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "Skipping server startup" in result.output assert "telemetry" not in runner.invoke(run_server, ["--help"]).output assert live_proxy_samples.exists() @@ -792,9 +744,7 @@ class TestProxyInitializationHelpers: result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert not live_proxy_samples.exists() @@ -831,9 +781,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_limit_concurrency_passed_to_uvicorn( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ): @@ -850,11 +798,7 @@ class TestProxyInitializationHelpers: KeyManagementSettings=MagicMock(), save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -864,9 +808,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.side_effect = lambda *a, **k: { "app": "litellm.proxy.proxy_server:app", @@ -874,20 +816,14 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--limit_concurrency", "250"] - ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + result = runner.invoke(run_server, ["--local", "--limit_concurrency", "250"]) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert mock_uvicorn_run.call_args.kwargs.get("limit_concurrency") == 250 mock_uvicorn_run.reset_mock() result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert "limit_concurrency" not in mock_uvicorn_run.call_args.kwargs @@ -905,7 +841,9 @@ class TestProxyInitializationHelpers: @patch("httpx.HTTPTransport.handle_request") @patch("atexit.register") @patch("subprocess.Popen") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above @patch( # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False ) @@ -943,9 +881,7 @@ class TestProxyInitializationHelpers: save_worker_config=MagicMock(), ) clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") } clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) with ( @@ -971,9 +907,7 @@ class TestProxyInitializationHelpers: run_server, ["--local", "--host", "127.0.0.1", "--port", "4000", "--prometheus_metrics_port", "4001"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_popen.assert_called_once() spawned = list(mock_popen.call_args.args[0]) assert spawned[1:3] == ["-m", "litellm.proxy.prometheus_metrics_server"] @@ -996,9 +930,7 @@ class TestProxyInitializationHelpers: mock_popen.reset_mock() mock_uvicorn_run.reset_mock() result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_popen.assert_not_called() mock_uvicorn_run.assert_called_once() @@ -1012,9 +944,7 @@ class TestProxyInitializationHelpers: mock_popen.assert_not_called() mock_uvicorn_run.assert_not_called() - result = runner.invoke( - run_server, ["--local", "--prometheus_metrics_port", "0"] - ) + result = runner.invoke(run_server, ["--local", "--prometheus_metrics_port", "0"]) assert result.exit_code == 2 assert "Invalid value for '--prometheus_metrics_port'" in result.output mock_popen.assert_not_called() @@ -1036,9 +966,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_db_timeout_settings_are_forwarded_to_pool_timeout( self, mock_should_update, @@ -1071,11 +999,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1086,9 +1010,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: ( @@ -1107,9 +1029,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] assert appended_params["connection_limit"] == 5 @@ -1164,9 +1084,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_db_connection_extra_params_forwarded_to_url( self, mock_should_update, @@ -1201,11 +1119,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1216,9 +1130,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1235,9 +1147,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] assert appended_params["connect_timeout"] == 15 @@ -1288,9 +1198,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_disable_prepared_statements_forwarded_to_url( self, mock_should_update, @@ -1322,11 +1230,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1337,9 +1241,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1356,9 +1258,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] if expect_pgbouncer: @@ -1380,9 +1280,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_disable_prepared_statements_env_var_forwarded_to_url( self, mock_should_update, @@ -1430,9 +1328,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1449,9 +1345,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" appended_params = mock_append_query_params.call_args.args[1] if expect_pgbouncer: assert appended_params["pgbouncer"] == "true", appended_params @@ -1461,9 +1355,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_malformed_disable_prepared_statements_env_var_is_rejected_even_when_config_enables_it( self, mock_should_update, @@ -1509,9 +1401,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1535,9 +1425,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_proxy_default_api_version_uses_azure_default( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ): @@ -1554,11 +1442,7 @@ class TestProxyInitializationHelpers: KeyManagementSettings=MagicMock(), save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1568,9 +1452,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1578,9 +1460,7 @@ class TestProxyInitializationHelpers: "port": 8000, } result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_proxy_module.save_worker_config.assert_called_once() call_kwargs = mock_proxy_module.save_worker_config.call_args[1] assert call_kwargs["api_version"] == litellm.AZURE_DEFAULT_API_VERSION @@ -1588,12 +1468,8 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) - def test_keepalive_timeout_flag( - self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run - ): + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_keepalive_timeout_flag(self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the keepalive_timeout flag is properly passed to uvicorn""" from click.testing import CliRunner @@ -1610,11 +1486,7 @@ class TestProxyInitializationHelpers: # DB-setup block (un-timeout'd `subprocess.run(["prisma"])` + # migrate-deploy retry loop) — same isolation every other run_server # test in this file uses. - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1629,9 +1501,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers.is_port_in_use", return_value=False, @@ -1664,12 +1534,8 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) - def test_timeout_worker_healthcheck_flag( - self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run - ): + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_timeout_worker_healthcheck_flag(self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the --timeout_worker_healthcheck flag is threaded through to the uvicorn init helper.""" from click.testing import CliRunner @@ -1686,11 +1552,7 @@ class TestProxyInitializationHelpers: # DB-setup block (un-timeout'd `subprocess.run(["prisma"])` + # migrate-deploy retry loop) — same isolation every other run_server # test in this file uses. - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), @@ -1705,9 +1567,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers.is_port_in_use", return_value=False, @@ -1719,9 +1579,7 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--timeout_worker_healthcheck", "15"] - ) + result = runner.invoke(run_server, ["--local", "--timeout_worker_healthcheck", "15"]) assert result.exit_code == 0 mock_get_args.assert_called_once_with( @@ -1736,9 +1594,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_max_requests_before_restart_flag( - self, mock_setup_db, mock_print, mock_uvicorn_run - ): + def test_max_requests_before_restart_flag(self, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests""" from click.testing import CliRunner @@ -1751,11 +1607,7 @@ class TestProxyInitializationHelpers: mock_key_mgmt = MagicMock() mock_save_worker_config = MagicMock() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict( os.environ, @@ -1773,9 +1625,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1783,13 +1633,9 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--max_requests_before_restart", "123"] - ) + result = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() # Check that uvicorn.run was called with limit_max_requests parameter @@ -1799,9 +1645,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_max_requests_before_restart_jitter_flag( - self, mock_setup_db, mock_print, mock_uvicorn_run - ): + def test_max_requests_before_restart_jitter_flag(self, mock_setup_db, mock_print, mock_uvicorn_run): """--max_requests_before_restart_jitter maps to uvicorn limit_max_requests_jitter""" from click.testing import CliRunner @@ -1812,11 +1656,7 @@ class TestProxyInitializationHelpers: pass runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch("uvicorn.Config", _NewUvicornConfig), @@ -1831,9 +1671,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1852,9 +1690,7 @@ class TestProxyInitializationHelpers: ], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() call_args = mock_uvicorn_run.call_args assert call_args[1]["limit_max_requests"] == 1000 @@ -1873,11 +1709,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1891,9 +1723,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1913,9 +1743,7 @@ class TestProxyInitializationHelpers: ], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_not_called() mock_run_gunicorn.assert_called_once() g_kwargs = mock_run_gunicorn.call_args[1] @@ -1961,20 +1789,14 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_single_worker_max_requests_warning_printed_on_uvicorn( - self, mock_setup_db, mock_uvicorn_run - ): + def test_single_worker_max_requests_warning_printed_on_uvicorn(self, mock_setup_db, mock_uvicorn_run): """uvicorn single-worker startup prints the does-not-recycle warning""" from click.testing import CliRunner from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -1988,9 +1810,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1998,13 +1818,9 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--max_requests_before_restart", "123"] - ) + result = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "--max_requests_before_restart 123 with --num_workers 1" in result.output assert "uvicorn exits the process" in result.output @@ -2021,11 +1837,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -2039,9 +1851,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2054,12 +1864,8 @@ class TestProxyInitializationHelpers: ["--local", "--num_workers", "2", "--max_requests_before_restart", "123"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" - assert not any( - "does not recycle a worker" in str(c) for c in mock_print.call_args_list - ) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert not any("does not recycle a worker" in str(c) for c in mock_print.call_args_list) @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server") @patch("uvicorn.run") @@ -2073,11 +1879,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -2091,9 +1893,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2106,9 +1906,7 @@ class TestProxyInitializationHelpers: ["--local", "--run_gunicorn", "--max_requests_before_restart", "123"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "--max_requests_before_restart 123 with --num_workers 1" in result.output assert "gunicorn accepts no connections" in result.output @@ -2125,11 +1923,7 @@ class TestProxyInitializationHelpers: from litellm.proxy.proxy_cli import run_server runner = CliRunner() - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict(os.environ, clean_env, clear=True), patch.dict( @@ -2143,9 +1937,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2158,13 +1950,9 @@ class TestProxyInitializationHelpers: ["--local", "--run_hypercorn", "--max_requests_before_restart", "5"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_init_hypercorn.assert_called_once() - assert not any( - "does not recycle a worker" in str(c) for c in mock_print.call_args_list - ) + assert not any("does not recycle a worker" in str(c) for c in mock_print.call_args_list) @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_options_include_max_requests_jitter(self): @@ -2244,6 +2032,89 @@ class TestProxyInitializationHelpers: assert worker.config.h11_max_incomplete_event_size == max_request_header_size assert captured["options"]["graceful_timeout"] == 8 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + @staticmethod + def _gunicorn_application(**listener: object): + captured: Final[list] = [] + + def capture_run(application) -> None: + captured.append(application) + + with ( + patch("gunicorn.app.base.BaseApplication.run", capture_run), + patch( # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership + "litellm.proxy.proxy_cli.start_query_engine_reaper" + ), + ): + ProxyInitializationHelpers._run_gunicorn_server( + host="0.0.0.0", + port=4012, + app=MagicMock(), + num_workers=2, + ssl_certfile_path=None, + ssl_keyfile_path=None, + **listener, + ) + return captured[0] + + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_serving_clients_directly_keeps_its_host_port_and_its_own_trusted_proxies(self): + pytest.importorskip("gunicorn") + import gunicorn.config + + config: Final = self._gunicorn_application().cfg + assert (config.address, config.forwarded_allow_ips) == ( + [("0.0.0.0", 4012)], + gunicorn.config.Config().forwarded_allow_ips, + ) + + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_listens_on_the_socket_it_is_handed_instead_of_its_host_and_port(self): + pytest.importorskip("gunicorn") + import socket + + import gunicorn.glogging + import gunicorn.sock + + with socket.socket() as handed: + handed.bind(("127.0.0.1", 0)) + handed.listen() + config: Final = self._gunicorn_application(private_listener=InheritedFd(os.dup(handed.fileno()))).cfg + listeners: Final = gunicorn.sock.create_sockets(config, gunicorn.glogging.Logger(config)) + try: + assert [listener.getsockname() for listener in listeners] == [handed.getsockname()] + finally: + gunicorn.sock.close_sockets(listeners, unlink=False) + + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") + def test_gunicorn_workers_on_a_private_listener_ignore_a_forwarded_client(self, monkeypatch): + pytest.importorskip("gunicorn") + import gunicorn.config + import gunicorn.glogging + + for name in ("uvicorn.error", "uvicorn.access"): + uvicorn_logger = logging.getLogger(name) + for attr in ("handlers", "level", "propagate"): + monkeypatch.setattr(uvicorn_logger, attr, getattr(uvicorn_logger, attr)) + application: Final = self._gunicorn_application(private_listener=AssignedPort(45556)) + worker: Final = application.options["worker_class"]( + 0, os.getpid(), [], MagicMock(), 30, application.cfg, gunicorn.glogging.Logger(gunicorn.config.Config()) + ) + seen: Final[list[str]] = [] + + async def probe(scope, receive, send) -> None: + seen.append(scope["client"][0]) + + worker.config.app = probe + worker.config.load() + request_from_loopback: Final = { + "type": "http", + "client": ("127.0.0.1", 51234), + "headers": [(b"x-forwarded-for", b"203.0.113.9")], + } + asyncio.run(worker.config.loaded_app(request_from_loopback, AsyncMock(), AsyncMock())) + + assert seen == ["127.0.0.1"] + @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation): """preload forks workers from the master, so native routes are forbidden there first""" @@ -2370,7 +2241,9 @@ class TestProxyInitializationHelpers: with patch.dict(os.environ, test_env_special): result = construct_database_url_from_env_vars() - expected_url = "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + expected_url = ( + "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + ) assert result == expected_url # Test without password (should still work) @@ -2515,12 +2388,8 @@ class TestQueryEngineReaperWiring: }, ), patch("uvicorn.run") as mock_uvicorn_run, - patch( - "litellm.proxy.proxy_cli.start_query_engine_reaper" - ) as mock_start_reaper, - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.start_query_engine_reaper") as mock_start_reaper, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2531,17 +2400,13 @@ class TestQueryEngineReaperWiring: return result, mock_uvicorn_run, mock_start_reaper def test_multi_worker_uvicorn_starts_reaper(self): - result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( - ["--local", "--num_workers", "2"] - ) + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(["--local", "--num_workers", "2"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() mock_start_reaper.assert_called_once() def test_single_worker_uvicorn_does_not_start_reaper(self): - result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( - ["--local", "--num_workers", "1"] - ) + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(["--local", "--num_workers", "1"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() mock_start_reaper.assert_not_called() @@ -2566,9 +2431,7 @@ class TestQueryEngineReaperWiring: with ( patch("gunicorn.app.base.BaseApplication.run"), - patch( - "litellm.proxy.proxy_cli.start_query_engine_reaper" - ) as mock_start_reaper, + patch("litellm.proxy.proxy_cli.start_query_engine_reaper") as mock_start_reaper, ): ProxyInitializationHelpers._run_gunicorn_server( host="127.0.0.1", @@ -2616,11 +2479,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" with ( @@ -2632,9 +2491,7 @@ class TestRunServerDbSetup: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2665,10 +2522,18 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_with(use_migrate=False, use_v2_resolver=True) @patch("atexit.register") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.check_migration.check_prisma_schema_diff" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above def test_migrations_run_when_the_prisma_cli_is_not_on_path( self, mock_should_update_schema, @@ -2692,11 +2557,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" clean_env["PATH"] = str(empty_bin) @@ -2856,11 +2717,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" with ( @@ -2872,9 +2729,7 @@ class TestRunServerDbSetup: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2917,11 +2772,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db" with ( @@ -2935,9 +2786,7 @@ class TestRunServerDbSetup: ), ): with pytest.raises(SystemExit) as exc_info: - run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) assert exc_info.value.code == 1 mock_setup_database.assert_not_called() err: Final = capsys.readouterr().err @@ -2979,11 +2828,7 @@ class TestRunServerDbSetup: save_worker_config=MagicMock(), ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" clean_env["USE_V2_MIGRATION_RESOLVER"] = "true" @@ -2997,9 +2842,7 @@ class TestRunServerDbSetup: }, ), ): - run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out @@ -3042,9 +2885,7 @@ class TestRunServerDbSetup: {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, ), ): - run_server.main( - ["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False) mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated and has no effect" in capsys.readouterr().out @@ -3181,12 +3022,7 @@ class TestRunServerDbSetup: def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected): from litellm.proxy.proxy_cli import resolve_v2_migration_resolver - assert ( - resolve_v2_migration_resolver( - use_legacy_flag=use_legacy_flag, env_value=env_value - ) - is expected - ) + assert resolve_v2_migration_resolver(use_legacy_flag=use_legacy_flag, env_value=env_value) is expected def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self): from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli @@ -3227,9 +3063,7 @@ class TestRunServerDbSetup: ) clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") } clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" @@ -3310,9 +3144,7 @@ class TestRunServerDbSetup: "sys.modules", {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, outcome as exc_info, ): mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000} @@ -3364,11 +3196,7 @@ class TestWorkerStartupHooks: "LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_dummy_hook", } # Remove DATABASE_URL to avoid real DB setup - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3391,11 +3219,7 @@ class TestWorkerStartupHooks: env_overrides = { "LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_dummy_async_hook", } - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3415,11 +3239,7 @@ class TestWorkerStartupHooks: env_overrides = { "LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_failing_hook", } - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3448,18 +3268,11 @@ class TestWorkerStartupHooks: from litellm.proxy.proxy_server import proxy_startup_event - hooks = ( - "tests.unit.proxy.test_proxy_cli:_dummy_hook," - "tests.unit.proxy.test_proxy_cli:_dummy_async_hook" - ) + hooks = "tests.unit.proxy.test_proxy_cli:_dummy_hook,tests.unit.proxy.test_proxy_cli:_dummy_async_hook" env_overrides = { "LITELLM_WORKER_STARTUP_HOOKS": hooks, } - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} clean_env.update(env_overrides) with patch.dict(os.environ, clean_env, clear=True): @@ -3605,13 +3418,9 @@ class TestPostgresStatementTimeoutOptions: import yaml config_path = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": 60}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": 60}})) - captured = _run_server_and_capture_urls( - str(config_path), direct_url="postgresql://t:t@localhost:5432/t" - ) + captured = _run_server_and_capture_urls(str(config_path), direct_url="postgresql://t:t@localhost:5432/t") assert "options" in urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query) assert "options" not in urlparse.parse_qs(urlparse.urlparse(captured["DIRECT_URL"]).query) @@ -3783,9 +3592,7 @@ class TestReadReplicaConnectionParams: import yaml config_path: Final = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})) _run_server_and_capture_urls( str(config_path), @@ -3805,9 +3612,7 @@ class TestReadReplicaConnectionParams: import yaml config_path: Final = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})) _run_server_and_capture_urls( str(config_path), @@ -4002,9 +3807,7 @@ class TestTokenAuthCliFlags: patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False), patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database"), patch("uvicorn.run"), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -4018,9 +3821,7 @@ class TestTokenAuthCliFlags: return result, database_url, toggle def test_azure_flag_assembles_a_token_bearing_database_url(self): - result, database_url, toggle = self._invoke_with_azure_host( - ["--local", "--azure_postgresql_auth"] - ) + result, database_url, toggle = self._invoke_with_azure_host(["--local", "--azure_postgresql_auth"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert database_url is not None @@ -4248,7 +4049,11 @@ def probe_logger() -> Iterator[logging.Logger]: def _log_config_source(kind: str, tmp_path: Path) -> str | dict[str, object]: - config: Final[dict[str, object]] = {"version": 1, "incremental": True, "loggers": {_PROBE_LOGGER: {"level": "ERROR"}}} + config: Final[dict[str, object]] = { + "version": 1, + "incremental": True, + "loggers": {_PROBE_LOGGER: {"level": "ERROR"}}, + } if kind == "dict": return config written: Final = tmp_path / f"log_config.{kind}" diff --git a/tests/unit/proxy/test_server_listener.py b/tests/unit/proxy/test_server_listener.py new file mode 100644 index 00000000000..ec2f7a44ad4 --- /dev/null +++ b/tests/unit/proxy/test_server_listener.py @@ -0,0 +1,28 @@ +import os +from typing import Final + +import pytest + +from litellm.proxy.server_listener import AssignedPort, InheritedFd, PrivateListener, gunicorn_overrides + + +def test_gunicorn_keeps_its_own_listener_and_trust_without_a_private_listener() -> None: + assert dict(gunicorn_overrides(None)) == {} + + +@pytest.mark.skipif(os.name == "nt", reason="gunicorn does not run on Windows") +@pytest.mark.parametrize( + ("listener", "address"), + [ + (InheritedFd(17), 17), + (AssignedPort(45556), ("127.0.0.1", 45556)), + ], +) +def test_gunicorn_serves_a_private_listener_and_trusts_no_peer_to_name_the_client( + listener: PrivateListener, address: int | tuple[str, int] +) -> None: + gunicorn_config: Final = pytest.importorskip("gunicorn.config").Config() + gunicorn_config.set("forwarded_allow_ips", "*") + for name, value in gunicorn_overrides(listener).items(): + gunicorn_config.set(name, value) + assert (gunicorn_config.address, gunicorn_config.forwarded_allow_ips) == ([address], []) diff --git a/tests/unit/rust_bridge/test_configuration.py b/tests/unit/rust_bridge/test_configuration.py index 8116cf98076..70223421471 100644 --- a/tests/unit/rust_bridge/test_configuration.py +++ b/tests/unit/rust_bridge/test_configuration.py @@ -106,8 +106,10 @@ def test_process_override_applies_when_environment_is_unset() -> None: assert configuration.rust_enabled() is True -@pytest.mark.parametrize("value", ("", " ", "sometimes", "2")) -def test_invalid_environment_value_is_ignored(monkeypatch: pytest.MonkeyPatch, value: str) -> None: +@pytest.mark.parametrize("value", ("", " ", "sometimes", "2", "enabled")) +def test_blank_or_invalid_environment_value_leaves_the_switch_unset( + monkeypatch: pytest.MonkeyPatch, value: str +) -> None: monkeypatch.setenv("LITELLM_RUST", value) assert configuration.rust_enabled() is False @@ -116,6 +118,13 @@ def test_invalid_environment_value_is_ignored(monkeypatch: pytest.MonkeyPatch, v assert configuration.rust_enabled() is True +def test_environment_is_read_on_every_decision(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_RUST", "1") + assert configuration.rust_enabled() is True + monkeypatch.setenv("LITELLM_RUST", "0") + assert configuration.rust_enabled() is False + + def test_process_override_and_reset_apply_to_existing_threads() -> None: with ThreadPoolExecutor(max_workers=1) as executor: assert executor.submit(configuration.rust_enabled).result() is False