scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200
litellm/tests/unit/proxy/test_proxy_cli.py
yucheng-berri a53ea5bd77
fix(proxy): keep every startup line JSON when json_logs is on (#47035) 
* fix(proxy): keep every startup line JSON when json_logs is on

Boot-time messages the proxy printed straight to stdout (banner, feedback
box, model list, cache and callback notices, policy table, migration
notices, CLI launch notes) and the first records of uvicorn- or
granian-spawned workers came out as plain text while json_logs was on,
breaking log pipelines that parse every line as JSON.

- print_startup_message() prints unchanged in plain mode and, with
  json_logs on, emits one INFO record on the proxy logger (ANSI
  stripped) that skips the logger and handler level checks like the
  print it replaces; show_banner() routes to it in JSON mode only
- the CLI exports JSON_LOGS=true when the config sets json_logs so
  spawned workers log JSON from their first record
- the --enforce_prisma_migration_check deprecation notice prints after
  the config is read so a config-only json_logs covers it

Resolves LIT-1624

* fix(proxy): keep show_banner() SDK-free, pick the JSON banner in initialize()

Since #46803 the `lite` shell calls show_banner() against an uninitialized
litellm package, and reading litellm.json_logs inside it initialized the whole
SDK at the banner (0.14 s to 0.83 s for a bare `lite` here, the seconds that PR
removed on small VMs). show_banner() is back to its main version; the proxy's
initialize() now emits the JSON banner record itself and calls show_banner()
only in plain mode. The banner unit tests move to the initialize() call site.

* fix(proxy): keep the JSON banner when an operator --log_config names the proxy logger

uvicorn applies --log_config after proxy_cli has switched JSON logging on and before
initialize() switches it on again through load_config(). A logging config that names
"LiteLLM Proxy" removes that logger's handlers in the window between the two, which is where
initialize() emits the banner and the feedback box, so at one worker both went nowhere;
two-worker children import litellm after uvicorn's logging setup and were not affected.
print_startup_message() now falls back to a one-off JSON handler, on the stream the JSON
handler uses, when the proxy logger has no handler of its own.

Live, json_logs on plus a log_config naming the proxy logger, one worker: main prints the
banner and the feedback box as plaintext, e2303a3925 prints neither, this head prints both as
one JSON record each. New unit test fails on e2303a3925; new integration cell
test_json_logs_with_operator_log_config_keeps_every_line_json passes at one and two workers.

* fix(proxy): keep a long startup listing whole in its JSON record

print_startup_message ran its record through the proxy logger's filters, so
in JSON mode the required-migration SQL list and the policy list were
middle-truncated at MAX_STRING_LENGTH_STDOUT_LOG (4096 by default) with the
litellm_truncated marker, where the prints they replace showed them whole.
That cap bounds echoed request payloads; a startup message is operator text,
so the record now gets the uncapped processing of a DEBUG line: secret
redaction and base64 collapse still apply, the text cap does not.

Red on 41f92b745a, green here:
- tests/unit/test_logging.py: a listing longer than the cap comes out whole;
  a credential in a startup message is still redacted
- tests/integration/configuration/test_json_logs_startup.py: a 60-policy
  config boots with json_logs on and the policy listing record (9762 chars)
  carries every policy with no truncation marker (4096 chars, marker, on
  the unfixed tree)

* fix(proxy): keep the --enforce_prisma_migration_check notice on --health and --test

This PR moved the deprecation notice after the config is read, so that a
json_logs set only in the config applies to it; --health and --test return
before that point and stopped printing it, a visible change main did not
have. The notice is now also emitted before those two returns, through the
same helper, so a deployment running `litellm --health
--enforce_prisma_migration_check` still learns the flag is dead, as on main.
A normal server start emits it once, after the config, as in the previous
commit.

Red on 2d7a2d6844, green here: tests/unit/proxy/test_proxy_cli.py asserts
the notice on --health and on --test.

* fix(proxy): keep the JSON startup records on stdout under an operator handler on the proxy logger

An operator --log_config that gives "LiteLLM Proxy" a handler of its own, a NullHandler
that silences it or a plain WARNING-level handler on stderr, is in force in the window
where the banner, the feedback box and the cache notice are emitted: uvicorn applies it
after the CLI process switched JSON on and before initialize() re-installs the JSON
handler through load_config(). 41f92b745a sent the startup records through the proxy
logger's handlers and fell back to a JSON stdout handler only when there were none, so
under such a config the records landed in the operator's handler, as plaintext on stderr,
or vanished, where the prints they replace reached stdout whatever the logging config
said.

print_startup_message() now writes every startup record through the proxy's own JSON
handler on stdout, filters included, and never through the proxy logger's handlers; the
operator's handler gets nothing from a startup message, as it got nothing from the print.
The normal JSON handler is the same class, formatter and filters, so the records a
pipeline reads are unchanged when no operator config is in play.

Tests: tests/unit/test_logging.py gets the two handler shapes (red on f6f5957d16: the
record went to stderr as plaintext or nowhere); tests/integration/configuration/
test_json_logs_startup.py gets a silencing operator --log_config cell at 1 and 2 workers
(red on f6f5957d16: banner, feedback box and cache notice missing; green here).

* fix(proxy): print the --enforce_prisma_migration_check notice when the config fails to load

f6f5957d16 printed the deprecation notice after the config is read, so a json_logs set
only in the config makes it a JSON record. A boot whose config file is missing or does
not parse then exited before that point and never printed it, where main printed it
before the error (10-20 gauntlet finding).

The notice is printed where main printed it, unless a config file can still switch
json_logs on: then it is held until the config is read and printed right after the
json_logs switch, and a click close callback prints it on the way out if the boot fails
first. functools.cache keeps it to one print either way.

Tests: tests/unit/proxy/test_proxy_cli.py covers a failed config read (red on
75043aed83), a plain config printing it once and a config-only json_logs making it one
JSON record; tests/integration/configuration/test_json_logs_startup.py adds the JSON cell
and a missing and an unparseable config cell (red on 75043aed83, green here), through
refused_boot_log, which now takes extra CLI arguments.

* fix(proxy): keep startup records as printed instead of running secret redaction

In JSON mode the startup messages ran secret redaction, which the prints they replace never
had. Its sk- key pattern has no left boundary, so a model named risk-assessment-gpt4o was
logged as riREDACTED in the model list (10-21 gauntlet finding). The records now skip
redaction like the prints did; the processed record still counts as scrubbed, so the
handler's secret filter and the formatter leave it alone. The stdout cap stays off and base64
collapse stays on, as before.

Tests: tests/unit/test_logging.py checks the model name comes out whole (red on 7d17f126e5),
replacing the test that asserted startup records were redacted;
tests/integration/configuration/test_json_logs_startup.py adds a boot with a
risk-assessment-<id> model whose model list record carries the name (red on 7d17f126e5).

* refactor(proxy): meet the AGENTS.md rules Greptile flagged on the startup-message change

The policy listing builds its rows as tuples in two helpers instead of appending to a list;
plain output is unchanged (same bytes as before across 3000 random policy and attachment
configs). The explanatory comments added with this change are gone, and the docstring of
print_startup_message now says the record is whole and unredacted.

Tests: the migration-diff JSON test no longer runs the fake prisma CLI; it checks the
skipped-diff notice in-process, and the fixture extension it needed is reverted. The
json_logs export and deprecation-notice CLI tests run the real proxy module and a real
config file instead of a mocked ProxyConfig and proxy module, and the --health/--test
notice tests no longer patch our helpers. Each still fails when the export, the hold or
call_on_close is removed. New test parameters are typed.

---------

Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
2026-10-27 14:49:07 -07:00

4075 lines
156 KiB
Python
Raw Blame History

import inspect
import json
import logging
import os
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import click
import fastapi
import httpx
import pytest
import builtins
import runpy
import sys
import types
import urllib.parse as urlparse
import uvicorn
import yaml
from uvicorn.config import LOOP_FACTORIES
from uvicorn.importer import import_from_string
import litellm
from litellm._logging import turn_on_json
from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
from litellm.proxy import proxy_cli
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT
@pytest.fixture(autouse=True)
def fork_reservation():
"""Reserving is irreversible: it would forbid native routes in this pytest worker for good"""
with patch( # test-quality-ok: process-global native state, a real reservation would poison every later test in the worker
"litellm.rust_bridge.fork_guard.reserve_process_for_forking"
) as reserve:
yield reserve
@pytest.mark.xdist_group("proxy_cli")
class TestProxyInitializationHelpers:
@patch("importlib.metadata.version")
@patch("click.echo")
def test_echo_litellm_version(self, mock_echo, mock_version):
# Setup
mock_version.return_value = "1.0.0"
# Execute
ProxyInitializationHelpers._echo_litellm_version()
# Assert
mock_version.assert_called_once_with("litellm")
mock_echo.assert_called_once_with("\nLiteLLM: Current Version = 1.0.0\n")
@patch("httpx.get")
@patch("builtins.print")
@patch("json.dumps")
def test_run_health_check(self, mock_dumps, mock_print, mock_get):
# Setup
mock_response = MagicMock()
mock_response.json.return_value = {"status": "healthy"}
mock_get.return_value = mock_response
mock_dumps.return_value = '{"status": "healthy"}'
# Execute
ProxyInitializationHelpers._run_health_check("localhost", 8000)
# Assert
mock_get.assert_called_once_with(url="http://localhost:8000/health")
mock_response.json.assert_called_once()
mock_dumps.assert_called_once_with({"status": "healthy"}, indent=4)
@patch("openai.OpenAI")
@patch("click.echo")
@patch("builtins.print")
def test_run_test_chat_completion(self, mock_print, mock_echo, mock_openai):
# Setup
mock_client = MagicMock()
mock_openai.return_value = mock_client
mock_response = MagicMock()
mock_client.chat.completions.create.return_value = mock_response
mock_stream_response = MagicMock()
mock_stream_response.__iter__.return_value = [MagicMock(), MagicMock()]
mock_client.chat.completions.create.side_effect = [
mock_response,
mock_stream_response,
]
# Execute
with pytest.raises(ValueError, match="Invalid test value"):
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"
)
# Assert
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
)
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"
)
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
)
# 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)
assert "version" in args["log_config"]
assert "formatters" in args["log_config"]
# Test with keepalive_timeout
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
)
assert args["log_config"] == "log_config.json"
assert args["timeout_keep_alive"] == 120
class _FakeUvicornConfig:
def __init__(self, timeout_worker_healthcheck=None):
pass
with patch("uvicorn.Config", _FakeUvicornConfig):
args = ProxyInitializationHelpers._get_default_unvicorn_init_args(
"localhost", 8000, timeout_worker_healthcheck=15
)
assert args["timeout_worker_healthcheck"] == 15
@pytest.mark.parametrize(
("raw_timeout", "expected"),
[("7.5", 7.5), ("0", 0.0), (None, DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT)],
)
def test_uvicorn_graceful_shutdown_is_bounded_by_graceful_shutdown_timeout(
self, monkeypatch, raw_timeout, expected
):
if raw_timeout is None:
monkeypatch.delenv("GRACEFUL_SHUTDOWN_TIMEOUT", raising=False)
else:
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", raw_timeout)
args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000)
assert uvicorn.Config(**args).timeout_graceful_shutdown == expected
@staticmethod
def _uvicorn_access_info_enabled(args: dict) -> bool:
import logging
loggers = tuple(logging.getLogger(n) for n in ("uvicorn", "uvicorn.error", "uvicorn.access", "uvicorn.asgi"))
saved = tuple((lg, lg.handlers[:], lg.level, lg.propagate) for lg in loggers)
try:
uvicorn.Config(**args).configure_logging()
return logging.getLogger("uvicorn.access").isEnabledFor(logging.INFO)
finally:
for lg, handlers, level, propagate in saved:
lg.handlers[:] = handlers
lg.setLevel(level)
lg.propagate = propagate
def test_litellm_log_error_silences_uvicorn_info_lines(self, monkeypatch):
monkeypatch.setenv("LITELLM_LOG", "ERROR")
args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000)
assert "log_config" not in args
assert self._uvicorn_access_info_enabled(args) is False
def test_unset_litellm_log_keeps_uvicorn_default_info_lines(self, monkeypatch):
monkeypatch.delenv("LITELLM_LOG", raising=False)
args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000)
assert "log_level" not in args
assert self._uvicorn_access_info_enabled(args) is True
def test_installed_uvicorn_supports_worker_flags(self):
params = inspect.signature(uvicorn.Config.__init__).parameters
assert "timeout_worker_healthcheck" in params
assert "limit_max_requests_jitter" in params
args = ProxyInitializationHelpers._get_default_unvicorn_init_args(
"localhost", 8000, timeout_worker_healthcheck=30
)
assert args["timeout_worker_healthcheck"] == 30
def test_get_reload_options_no_config_still_watches_env(self):
opts = ProxyInitializationHelpers._get_reload_options(None)
assert opts["reload"] is True
assert opts["reload_dirs"] == [os.path.abspath(os.getcwd())]
assert opts["reload_includes"] == ["*.py", ".env"]
def test_get_reload_options_with_config_in_cwd(self, tmp_path, monkeypatch):
config_file = tmp_path / "config.yaml"
config_file.write_text("model_list: []\n")
monkeypatch.chdir(tmp_path)
opts = ProxyInitializationHelpers._get_reload_options("config.yaml")
assert opts["reload"] is True
assert opts["reload_dirs"] == [str(tmp_path)]
assert opts["reload_includes"] == ["*.py", ".env", "config.yaml"]
def test_get_reload_options_with_config_outside_cwd(self, tmp_path, monkeypatch):
cwd_dir = tmp_path / "work"
cwd_dir.mkdir()
elsewhere = tmp_path / "configs"
elsewhere.mkdir()
config_file = elsewhere / "proxy.yaml"
config_file.write_text("model_list: []\n")
monkeypatch.chdir(cwd_dir)
opts = ProxyInitializationHelpers._get_reload_options(str(config_file))
assert opts["reload"] is True
assert opts["reload_dirs"] == [str(cwd_dir), str(elsewhere)]
assert opts["reload_includes"] == ["*.py", ".env", "proxy.yaml"]
def test_patch_statreload_extra_paths_yields_config_and_py(self, tmp_path):
from pathlib import Path
from uvicorn.supervisors.statreload import StatReload
if hasattr(StatReload, "_litellm_patched_config_paths"):
StatReload._litellm_patched_config_paths.clear()
config_file = tmp_path / "config.yaml"
config_file.write_text("model_list: []\n")
py_file = tmp_path / "module.py"
py_file.write_text("x = 1\n")
applied = ProxyInitializationHelpers._patch_statreload_extra_paths(
[str(config_file)]
)
assert applied is True
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
assert py_file.resolve() in yielded_paths
def test_patch_statreload_extra_paths_yields_env(self, tmp_path):
from pathlib import Path
from uvicorn.supervisors.statreload import StatReload
if hasattr(StatReload, "_litellm_patched_config_paths"):
StatReload._litellm_patched_config_paths.clear()
env_file = tmp_path / ".env"
env_file.write_text("FOO=bar\n")
applied = ProxyInitializationHelpers._patch_statreload_extra_paths(
[str(env_file)]
)
assert applied is True
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
def test_patch_statreload_extra_paths_skips_falsy(self, tmp_path):
from uvicorn.supervisors.statreload import StatReload
if hasattr(StatReload, "_litellm_patched_config_paths"):
StatReload._litellm_patched_config_paths.clear()
assert ProxyInitializationHelpers._patch_statreload_extra_paths([]) 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
from uvicorn.supervisors.statreload import StatReload
if hasattr(StatReload, "_litellm_patched_config_paths"):
StatReload._litellm_patched_config_paths.clear()
config_file = tmp_path / "config.yaml"
config_file.write_text("model_list: []\n")
py_file = tmp_path / "only.py"
py_file.write_text("x = 1\n")
for _ in range(3):
ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)])
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
):
from pathlib import Path
from uvicorn.supervisors.statreload import StatReload
if hasattr(StatReload, "_litellm_patched_config_paths"):
StatReload._litellm_patched_config_paths.clear()
monkeypatch.delenv("LITELLM_DEV_ENV_HOT_RELOAD", raising=False)
config_file = tmp_path / "config.yaml"
config_file.write_text("model_list: []\n")
env_file = tmp_path / ".env"
env_file.write_text("FOO=bar\n")
monkeypatch.chdir(tmp_path)
uvicorn_args: dict = {}
with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning:
ProxyInitializationHelpers._configure_dev_reload(
uvicorn_args, str(config_file)
)
assert os.environ["LITELLM_DEV_ENV_HOT_RELOAD"] == "True"
assert uvicorn_args["reload"] is True
assert ".env" in uvicorn_args["reload_includes"]
mock_warning.assert_called_once()
warning_text = mock_warning.call_args.args[0].lower()
assert "override" in warning_text
assert ".env" in warning_text
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
def test_dev_env_hot_reload_enabled_reads_flag(self, monkeypatch):
from litellm._process_defaults import _dev_env_hot_reload_enabled
monkeypatch.setenv("LITELLM_DEV_ENV_HOT_RELOAD", "True")
assert _dev_env_hot_reload_enabled() is True
monkeypatch.setenv("LITELLM_DEV_ENV_HOT_RELOAD", "false")
assert _dev_env_hot_reload_enabled() is False
monkeypatch.delenv("LITELLM_DEV_ENV_HOT_RELOAD", raising=False)
assert _dev_env_hot_reload_enabled() is False
@patch("asyncio.run")
@patch("builtins.print")
def test_init_hypercorn_server(self, mock_print, mock_asyncio_run):
# Setup
mock_app = MagicMock()
# Execute
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"
)
@patch("granian.Granian")
@patch("builtins.print")
def test_init_granian_server(self, mock_print, mock_granian_cls):
pytest.importorskip("granian")
mock_server = MagicMock()
mock_granian_cls.return_value = mock_server
fake_interfaces = SimpleNamespace(ASGI="asgi")
with patch("granian.constants.Interfaces", fake_interfaces):
ProxyInitializationHelpers._init_granian_server(
host="0.0.0.0",
port=4000,
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
max_requests_before_restart=None,
ciphers=None,
granian_runtime_threads=None,
)
mock_granian_cls.assert_called_once()
call_kwargs = mock_granian_cls.call_args.kwargs
assert call_kwargs["target"] == "litellm.proxy.proxy_server:app"
assert call_kwargs["address"] == "0.0.0.0"
assert call_kwargs["port"] == 4000
assert call_kwargs["workers"] == 2
assert call_kwargs["interface"] == "asgi"
assert call_kwargs["websockets"] is True
assert "runtime_threads" not in call_kwargs
mock_server.serve.assert_called_once()
@patch("granian.Granian")
@patch("builtins.print")
def test_init_granian_server_runtime_threads(self, mock_print, mock_granian_cls):
pytest.importorskip("granian")
mock_server = MagicMock()
mock_granian_cls.return_value = mock_server
fake_interfaces = SimpleNamespace(ASGI="asgi")
with patch("granian.constants.Interfaces", fake_interfaces):
ProxyInitializationHelpers._init_granian_server(
host="0.0.0.0",
port=4000,
num_workers=1,
ssl_certfile_path=None,
ssl_keyfile_path=None,
max_requests_before_restart=None,
ciphers=None,
granian_runtime_threads=4,
)
assert mock_granian_cls.call_args.kwargs["runtime_threads"] == 4
@patch("granian.Granian")
@patch("builtins.print")
def test_init_granian_server_ssl(self, mock_print, mock_granian_cls):
pytest.importorskip("granian")
mock_server = MagicMock()
mock_granian_cls.return_value = mock_server
fake_interfaces = SimpleNamespace(ASGI="asgi")
with patch("granian.constants.Interfaces", fake_interfaces):
ProxyInitializationHelpers._init_granian_server(
host="0.0.0.0",
port=4000,
num_workers=1,
ssl_certfile_path="/path/to/cert.pem",
ssl_keyfile_path="/path/to/key.pem",
max_requests_before_restart=None,
ciphers=None,
granian_runtime_threads=None,
)
call_kwargs = mock_granian_cls.call_args.kwargs
assert call_kwargs["ssl_cert"] == Path("/path/to/cert.pem")
assert call_kwargs["ssl_key"] == Path("/path/to/key.pem")
mock_server.serve.assert_called_once()
@patch("granian.Granian")
def test_init_granian_server_ssl_requires_cert_and_key(self, mock_granian_cls):
pytest.importorskip("granian")
fake_interfaces = SimpleNamespace(ASGI="asgi")
with patch("granian.constants.Interfaces", fake_interfaces):
with pytest.raises(click.ClickException, match="Both --ssl_certfile_path"):
ProxyInitializationHelpers._init_granian_server(
host="0.0.0.0",
port=4000,
num_workers=1,
ssl_certfile_path="/path/to/cert.pem",
ssl_keyfile_path=None,
max_requests_before_restart=None,
ciphers=None,
granian_runtime_threads=None,
)
mock_granian_cls.assert_not_called()
@patch("subprocess.Popen")
def test_run_ollama_serve(self, mock_popen):
# Execute
ProxyInitializationHelpers._run_ollama_serve()
# Assert
mock_popen.assert_called_once()
# Test exception handling
mock_popen.side_effect = Exception("Test exception")
ProxyInitializationHelpers._run_ollama_serve() # Should not raise
@patch("socket.socket")
def test_is_port_in_use(self, mock_socket):
# Setup for port in use
mock_socket_instance = MagicMock()
mock_socket_instance.connect_ex.return_value = 0
mock_socket.return_value.__enter__.return_value = mock_socket_instance
# Execute and Assert
assert ProxyInitializationHelpers._is_port_in_use(8000) is True
# Setup for port not in use
mock_socket_instance.connect_ex.return_value = 1
# Execute and Assert
assert ProxyInitializationHelpers._is_port_in_use(8000) is False
def test_get_loop_type(self):
# Test on Windows
with patch("sys.platform", "win32"):
assert ProxyInitializationHelpers._get_loop_type() is None
# Test on Linux
with patch("sys.platform", "linux"):
assert ProxyInitializationHelpers._get_loop_type() == "uvloop"
@pytest.mark.parametrize(("env_value", "expected"), [("asyncio", "asyncio"), ("uvloop", "uvloop"), ("", "uvloop")])
def test_get_loop_type_honors_litellm_event_loop_env(self, env_value: str, expected: str, monkeypatch) -> None:
monkeypatch.setattr("sys.platform", "linux")
monkeypatch.setenv("LITELLM_EVENT_LOOP", env_value)
assert ProxyInitializationHelpers._get_loop_type() == expected
def test_selected_loop_factory_imports_on_this_interpreter(self):
loop_type = ProxyInitializationHelpers._get_loop_type()
if loop_type is None:
pytest.skip("uvicorn picks the loop itself on this platform")
assert callable(import_from_string(LOOP_FACTORIES[loop_type]))
@patch.dict(os.environ, {}, clear=True)
def test_database_url_construction_with_special_characters(self):
# Setup environment variables with special characters that need escaping
test_env = {
"DATABASE_HOST": "localhost:5432",
"DATABASE_USERNAME": "user@with+special",
"DATABASE_PASSWORD": "test-password-special-chars",
"DATABASE_NAME": "db_name/test",
}
with patch.dict(os.environ, test_env):
# Call the relevant function - we'll need to extract the database URL construction logic
# This is simulating what happens in the run_server function when database_url is None
import urllib.parse
from litellm.proxy.proxy_cli import append_query_params
database_host = os.environ["DATABASE_HOST"]
database_username = os.environ["DATABASE_USERNAME"]
database_password = os.environ["DATABASE_PASSWORD"]
database_name = os.environ["DATABASE_NAME"]
# Test the URL encoding part
database_username_enc = urllib.parse.quote_plus(database_username)
database_password_enc = urllib.parse.quote_plus(database_password)
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}"
# 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"
assert database_url == expected_url
# Test appending query parameters
params = {"connection_limit": 10, "pool_timeout": 60}
modified_url = append_query_params(database_url, params)
assert "connection_limit=10" in modified_url
assert "pool_timeout=60" in modified_url
def test_append_query_params_handles_missing_url(self):
from litellm.proxy.proxy_cli import append_query_params
modified_url = append_query_params(None, {"connection_limit": 10})
assert modified_url == ""
@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
)
def test_skip_server_startup(
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run, tmp_path: Path
):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
# 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["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")
with (
patch.dict(
os.environ,
clean_env,
clear=True,
),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
# Prevent real import of proxy_server inside Click's
# isolation context (heavy side effects cause stream
# lifecycle issues with Click 8.2+)
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
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",
"host": "localhost",
"port": 8000,
}
# --- 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 "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}"
assert "Skipping server startup" in result.output
assert "telemetry" not in runner.invoke(run_server, ["--help"]).output
assert live_proxy_samples.exists()
# --- normal startup ---
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}"
mock_uvicorn_run.assert_called_once()
assert not live_proxy_samples.exists()
@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)
def test_script_boot_imports_the_package_proxy_server(
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run
):
package_proxy_server = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
sibling_proxy_server = types.ModuleType("proxy_server")
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(
"sys.modules",
{"proxy_server": sibling_proxy_server, "litellm.proxy.proxy_server": package_proxy_server},
),
patch.object(sys, "argv", ["proxy_cli.py", "--skip_server_startup"]),
patch.object(sys, "path", list(sys.path)),
pytest.raises(SystemExit) as exit_info,
):
runpy.run_path(proxy_cli.__file__, run_name="__main__")
assert exit_info.value.code == 0
package_proxy_server.save_worker_config.assert_called_once()
@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
)
def test_limit_concurrency_passed_to_uvicorn(
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run
):
"""--limit_concurrency must reach uvicorn.run so uvicorn sheds load with 503
past the cap; omitted values stay absent and non-positive values are rejected."""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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,
):
mock_get_args.side_effect = lambda *a, **k: {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"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}"
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}"
mock_uvicorn_run.assert_called_once()
assert "limit_concurrency" not in mock_uvicorn_run.call_args.kwargs
for invalid_value in ("0", "-1"):
mock_uvicorn_run.reset_mock()
result = runner.invoke(
run_server,
["--local", "--limit_concurrency", invalid_value],
)
assert result.exit_code == 2
assert "Invalid value for '--limit_concurrency'" in result.output
mock_uvicorn_run.assert_not_called()
@patch("uvicorn.run")
@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( # 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
)
def test_prometheus_metrics_port_starts_separate_metrics_process(
self,
mock_should_update,
mock_setup_db,
mock_popen,
mock_atexit_register,
mock_handle_request,
mock_uvicorn_run,
tmp_path,
):
"""--prometheus_metrics_port must spawn `python -m litellm.proxy.prometheus_metrics_server` on --host
with the shared multiproc dir, wait for its /metrics response, and only then start uvicorn. It must stay off by
default, refuse to share --port, and abort the proxy when the child dies before serving."""
import httpx
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_popen.return_value = MagicMock(pid=4242, **{"poll.return_value": None})
probed_urls: list[str] = []
def child_metrics(request: httpx.Request) -> httpx.Response:
probed_urls.append(str(request.url))
return httpx.Response(200, headers={"x-litellm-metrics-pid": "4242"}, content=b"")
mock_handle_request.side_effect = child_metrics
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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", "PROMETHEUS_METRICS_PORT")
}
clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path)
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
patch( # test-quality-ok: same isolation as the sibling CLI tests above
"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",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
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}"
mock_popen.assert_called_once()
spawned = list(mock_popen.call_args.args[0])
assert spawned[1:3] == ["-m", "litellm.proxy.prometheus_metrics_server"]
assert spawned[3:] == ["--host", "127.0.0.1", "--port", "4001", "--multiproc_dir", str(tmp_path)]
assert probed_urls == ["http://127.0.0.1:4001/metrics"]
assert "Serving Prometheus metrics on 127.0.0.1:4001/metrics (pid 4242)" in result.output
mock_uvicorn_run.assert_called_once()
mock_popen.reset_mock()
mock_uvicorn_run.reset_mock()
mock_popen.return_value = MagicMock(pid=4243, **{"poll.return_value": 1})
result = runner.invoke(
run_server,
["--local", "--port", "4000", "--prometheus_metrics_port", "4001"],
)
assert result.exit_code == 1, f"exit_code={result.exit_code}, output={result.output}"
assert "Prometheus metrics server exited with code 1 before serving 0.0.0.0:4001" in result.output
mock_uvicorn_run.assert_not_called()
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}"
mock_popen.assert_not_called()
mock_uvicorn_run.assert_called_once()
mock_uvicorn_run.reset_mock()
result = runner.invoke(
run_server,
["--local", "--port", "4000", "--prometheus_metrics_port", "4000"],
)
assert result.exit_code == 2
assert "--prometheus_metrics_port must differ from --port" in result.output
mock_popen.assert_not_called()
mock_uvicorn_run.assert_not_called()
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()
@pytest.mark.parametrize(
"timeout_config,expected_timeout",
[
({"database_connection_timeout": 30}, 30),
({"database_connection_pool_timeout": 45}, 45),
(
{
"database_connection_timeout": 30,
"database_connection_pool_timeout": 45,
},
30,
),
],
)
@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
)
def test_db_timeout_settings_are_forwarded_to_pool_timeout(
self,
mock_should_update,
mock_setup_db,
mock_atexit_register,
mock_subprocess_run,
timeout_config,
expected_timeout,
):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
return_value={
"general_settings": {
"database_url": "postgresql://test:test@localhost:5432/test",
"database_connection_pool_limit": 5,
**timeout_config,
}
}
)
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(
"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.append_query_params",
side_effect=lambda url, params: (
f"{url}?connection_limit={params['connection_limit']}&pool_timeout={params['pool_timeout']}"
),
) as mock_append_query_params,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--config", "test-config.yaml", "--skip_server_startup"],
)
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
assert appended_params["pool_timeout"] == expected_timeout
def test_build_db_connection_url_params_defaults(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(connection_limit=10, pool_timeout=60)
assert params == {"connection_limit": 10, "pool_timeout": 60}
def test_build_db_connection_url_params_omits_none_timeouts(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(
connection_limit=10,
pool_timeout=60,
connect_timeout=None,
socket_timeout=None,
)
assert "connect_timeout" not in params
assert "socket_timeout" not in params
def test_build_db_connection_url_params_includes_optional_timeouts(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(
connection_limit=10,
pool_timeout=60,
connect_timeout=15,
socket_timeout=120,
)
assert params["connect_timeout"] == 15
assert params["socket_timeout"] == 120
def test_build_db_connection_url_params_extras_override_defaults(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(
connection_limit=10,
pool_timeout=60,
extra_params={
"pgbouncer": "true",
"statement_cache_size": 0,
"pool_timeout": 5,
},
)
assert params["pgbouncer"] == "true"
assert params["statement_cache_size"] == 0
assert params["pool_timeout"] == 5
@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
)
def test_db_connection_extra_params_forwarded_to_url(
self,
mock_should_update,
mock_setup_db,
mock_atexit_register,
mock_subprocess_run,
):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
return_value={
"general_settings": {
"database_url": "postgresql://test:test@localhost:5432/test",
"database_connect_timeout": 15,
"database_socket_timeout": 120,
"database_extra_connection_params": {
"pgbouncer": "true",
"statement_cache_size": 0,
},
}
}
)
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(
"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.append_query_params",
side_effect=lambda url, params: str(url),
) as mock_append_query_params,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--config", "test-config.yaml", "--skip_server_startup"],
)
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
assert appended_params["socket_timeout"] == 120
assert appended_params["pgbouncer"] == "true"
assert appended_params["statement_cache_size"] == 0
def test_build_db_connection_url_params_disable_prepared_statements(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(
connection_limit=10,
pool_timeout=60,
disable_prepared_statements=True,
)
assert params["pgbouncer"] == "true"
def test_build_db_connection_url_params_no_pgbouncer_by_default(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(
connection_limit=10,
pool_timeout=60,
)
assert "pgbouncer" not in params
def test_build_db_connection_url_params_extra_pgbouncer_overrides_flag(self):
from litellm.proxy.proxy_cli import _build_db_connection_url_params
params = _build_db_connection_url_params(
connection_limit=10,
pool_timeout=60,
disable_prepared_statements=True,
extra_params={"pgbouncer": "false"},
)
assert params["pgbouncer"] == "false"
@pytest.mark.parametrize(
"config_value, expect_pgbouncer",
[
(True, True),
(False, False),
("true", True),
("false", False),
("not-a-bool", False),
],
)
@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
)
def test_disable_prepared_statements_forwarded_to_url(
self,
mock_should_update,
mock_setup_db,
mock_atexit_register,
mock_subprocess_run,
config_value,
expect_pgbouncer,
):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
return_value={
"general_settings": {
"database_url": "postgresql://test:test@localhost:5432/test",
"database_disable_prepared_statements": config_value,
}
}
)
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(
"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.append_query_params",
side_effect=lambda url, params: str(url),
) as mock_append_query_params,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--config", "test-config.yaml", "--skip_server_startup"],
)
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:
assert appended_params["pgbouncer"] == "true"
else:
assert "pgbouncer" not in appended_params
@pytest.mark.parametrize(
"env_value, config_value, expect_pgbouncer",
[
("true", None, True),
("1", None, True),
("false", None, False),
(None, None, False),
("true", False, True),
("false", True, True),
],
)
@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
)
def test_disable_prepared_statements_env_var_forwarded_to_url(
self,
mock_should_update,
mock_setup_db,
mock_atexit_register,
mock_subprocess_run,
env_value,
config_value,
expect_pgbouncer,
):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_subprocess_run.return_value = MagicMock(returncode=0)
general_settings = {"database_url": "postgresql://test:test@localhost:5432/test"}
if config_value is not None:
general_settings["database_disable_prepared_statements"] = config_value
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
return_value={"general_settings": general_settings}
)
clean_env = {
k: v
for k, v in os.environ.items()
if k not in ("DATABASE_URL", "DIRECT_URL", "DATABASE_DISABLE_PREPARED_STATEMENTS")
}
if env_value is not None:
clean_env["DATABASE_DISABLE_PREPARED_STATEMENTS"] = env_value
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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.append_query_params",
side_effect=lambda url, params: str(url),
) as mock_append_query_params,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--config", "test-config.yaml", "--skip_server_startup"],
)
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
else:
assert "pgbouncer" not in appended_params, appended_params
@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
)
def test_malformed_disable_prepared_statements_env_var_is_rejected_even_when_config_enables_it(
self,
mock_should_update,
mock_setup_db,
mock_atexit_register,
mock_subprocess_run,
):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
mock_proxy_module.ProxyConfig.return_value.get_config = AsyncMock(
return_value={
"general_settings": {
"database_url": "postgresql://test:test@localhost:5432/test",
"database_disable_prepared_statements": True,
}
}
)
clean_env = {
k: v
for k, v in os.environ.items()
if k not in ("DATABASE_URL", "DIRECT_URL", "DATABASE_DISABLE_PREPARED_STATEMENTS")
}
clean_env["DATABASE_DISABLE_PREPARED_STATEMENTS"] = "enabled"
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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.append_query_params",
side_effect=lambda url, params: str(url),
) as mock_append_query_params,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--config", "test-config.yaml", "--skip_server_startup"],
)
assert isinstance(result.exception, ValueError), f"exit_code={result.exit_code}, output={result.output}"
assert "DATABASE_DISABLE_PREPARED_STATEMENTS" in str(result.exception), result.exception
mock_append_query_params.assert_not_called()
@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
)
def test_proxy_default_api_version_uses_azure_default(
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run
):
"""Proxy default api_version should match litellm.AZURE_DEFAULT_API_VERSION for consistency."""
from click.testing import CliRunner
import litellm
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"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}"
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
@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
):
"""Test that the keepalive_timeout flag is properly passed to uvicorn"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
# Strip DATABASE_URL/DIRECT_URL so run_server doesn't enter the prisma
# 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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
),
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,
),
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
"timeout_keep_alive": 30,
}
result = runner.invoke(run_server, ["--local", "--keepalive_timeout", "30"])
assert result.exit_code == 0
mock_get_args.assert_called_once_with(
host="0.0.0.0",
port=4000,
log_config=None,
keepalive_timeout=30,
timeout_worker_healthcheck=None,
)
mock_uvicorn_run.assert_called_once()
# Check that the uvicorn.run was called with the timeout_keep_alive parameter
call_args = mock_uvicorn_run.call_args
assert call_args[1]["timeout_keep_alive"] == 30
@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
):
"""Test that the --timeout_worker_healthcheck flag is threaded through to the uvicorn init helper."""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
# Strip DATABASE_URL/DIRECT_URL so run_server doesn't enter the prisma
# 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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
),
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,
),
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server, ["--local", "--timeout_worker_healthcheck", "15"]
)
assert result.exit_code == 0
mock_get_args.assert_called_once_with(
host="0.0.0.0",
port=4000,
log_config=None,
keepalive_timeout=None,
timeout_worker_healthcheck=15,
)
@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
):
"""Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
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")
}
with (
patch.dict(
os.environ,
clean_env,
clear=True,
),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
},
),
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",
"host": "localhost",
"port": 8000,
}
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}"
mock_uvicorn_run.assert_called_once()
# Check that uvicorn.run was called with limit_max_requests parameter
call_args = mock_uvicorn_run.call_args
assert call_args[1]["limit_max_requests"] == 123
@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
):
"""--max_requests_before_restart_jitter maps to uvicorn limit_max_requests_jitter"""
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
class _NewUvicornConfig:
def __init__(self, limit_max_requests=None, limit_max_requests_jitter=0):
pass
runner = CliRunner()
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),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
[
"--local",
"--max_requests_before_restart",
"1000",
"--max_requests_before_restart_jitter",
"50",
],
)
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
assert call_args[1]["limit_max_requests_jitter"] == 50
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server")
@patch("uvicorn.run")
@patch("builtins.print")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
def test_run_gunicorn_passes_max_requests_jitter(
self, mock_setup_db, mock_print, mock_uvicorn_run, mock_run_gunicorn
):
"""--run_gunicorn threads jitter into _run_gunicorn_server, not uvicorn.run"""
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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
[
"--local",
"--run_gunicorn",
"--max_requests_before_restart",
"900",
"--max_requests_before_restart_jitter",
"75",
],
)
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]
assert g_kwargs["max_requests_before_restart"] == 900
assert g_kwargs["max_requests_before_restart_jitter"] == 75
@pytest.mark.parametrize("run_gunicorn", [False, True])
@pytest.mark.parametrize(
"num_workers,max_requests_before_restart",
[(1, None), (2, 5)],
)
def test_single_worker_max_requests_warning_returns_none(
self, num_workers, max_requests_before_restart, run_gunicorn
):
"""no warning when the flag is unset or more than one worker serves"""
warning = ProxyInitializationHelpers.single_worker_max_requests_warning(
num_workers=num_workers,
max_requests_before_restart=max_requests_before_restart,
run_gunicorn=run_gunicorn,
)
assert warning is None
def test_single_worker_max_requests_warning_uvicorn(self):
warning = ProxyInitializationHelpers.single_worker_max_requests_warning(
num_workers=1,
max_requests_before_restart=5,
run_gunicorn=False,
)
assert warning is not None
assert "--max_requests_before_restart 5" in warning
assert "--num_workers 1" in warning
assert "uvicorn exits the process" in warning
def test_single_worker_max_requests_warning_gunicorn(self):
warning = ProxyInitializationHelpers.single_worker_max_requests_warning(
num_workers=1,
max_requests_before_restart=5,
run_gunicorn=True,
)
assert warning is not None
assert "gunicorn accepts no connections" in warning
assert "uvicorn exits" not in warning
@patch("uvicorn.run")
@patch("builtins.print")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
def test_single_worker_max_requests_warning_printed_on_uvicorn(
self, mock_setup_db, mock_print, 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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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",
"host": "localhost",
"port": 8000,
}
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 any(
"--max_requests_before_restart 123 with --num_workers 1" in str(c)
and "uvicorn exits the process" in str(c)
for c in mock_print.call_args_list
)
@patch("litellm.proxy.proxy_cli.start_query_engine_reaper")
@patch("uvicorn.run")
@patch("builtins.print")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
def test_single_worker_max_requests_no_warning_with_two_workers(
self, mock_setup_db, mock_print, mock_uvicorn_run, mock_start_reaper
):
"""num_workers 2 with --max_requests_before_restart prints no 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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--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
)
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server")
@patch("uvicorn.run")
@patch("builtins.print")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
def test_single_worker_max_requests_warning_printed_on_gunicorn(
self, mock_setup_db, mock_print, mock_uvicorn_run, mock_run_gunicorn
):
"""--run_gunicorn single-worker startup prints the gunicorn warning variant"""
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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--run_gunicorn", "--max_requests_before_restart", "123"],
)
assert (
result.exit_code == 0
), f"exit_code={result.exit_code}, output={result.output}"
assert any(
"--max_requests_before_restart 123 with --num_workers 1" in str(c)
and "gunicorn accepts no connections" in str(c)
for c in mock_print.call_args_list
)
@patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.init_hypercorn_server")
@patch("uvicorn.run")
@patch("builtins.print")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
def test_single_worker_max_requests_no_warning_on_hypercorn(
self, mock_setup_db, mock_print, mock_uvicorn_run, mock_init_hypercorn
):
"""hypercorn ignores --max_requests_before_restart, so no warning prints"""
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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(
run_server,
["--local", "--run_hypercorn", "--max_requests_before_restart", "5"],
)
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
)
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_options_include_max_requests_jitter(self):
"""_run_gunicorn_server puts max_requests_jitter into the gunicorn options"""
pytest.importorskip("gunicorn")
captured: dict = {}
def capture_run(self):
captured["options"] = dict(self.options)
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="127.0.0.1",
port=4010,
app=MagicMock(),
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
max_requests_before_restart=1000,
max_requests_before_restart_jitter=50,
)
assert captured["options"]["max_requests"] == 1000
assert captured["options"]["max_requests_jitter"] == 50
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_workers_bound_graceful_shutdown_by_graceful_shutdown_timeout(self, monkeypatch):
pytest.importorskip("gunicorn")
import gunicorn.config
import gunicorn.glogging
from uvicorn.workers import UvicornWorker
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "7.5")
captured: dict = {}
def capture_run(self):
captured["options"] = dict(self.options)
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="127.0.0.1",
port=4011,
app=MagicMock(),
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
)
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))
gunicorn_cfg = gunicorn.config.Config()
worker = captured["options"]["worker_class"](
0, os.getpid(), [], MagicMock(), 30, gunicorn_cfg, gunicorn.glogging.Logger(gunicorn_cfg)
)
assert isinstance(worker, UvicornWorker)
assert worker.config.timeout_graceful_shutdown == 7.5
assert worker.config.loop == UvicornWorker.CONFIG_KWARGS["loop"]
assert worker.config.http == UvicornWorker.CONFIG_KWARGS["http"]
assert captured["options"]["graceful_timeout"] == 8 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
@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"""
pytest.importorskip("gunicorn")
reserved_before_run: list = []
def capture_run(self):
reserved_before_run.append(fork_reservation.call_args)
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="127.0.0.1",
port=4012,
app=MagicMock(),
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
)
assert [call.args for call in reserved_before_run] == [("the gunicorn master",)]
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_jitter_without_base_warns(self):
"""gunicorn path warns when jitter is set without --max_requests_before_restart"""
pytest.importorskip("gunicorn")
captured: dict = {}
def capture_run(self):
captured["options"] = dict(self.options)
with (
patch("gunicorn.app.base.BaseApplication.run", capture_run),
patch("builtins.print") as mock_print,
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="127.0.0.1",
port=4011,
app=MagicMock(),
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
max_requests_before_restart=None,
max_requests_before_restart_jitter=50,
)
assert "max_requests" not in captured["options"]
assert "max_requests_jitter" not in captured["options"]
assert any("has no effect" in str(c) for c in mock_print.call_args_list)
def test_apply_uvicorn_jitter_sets_arg_when_supported(self):
class _NewUvicornConfig:
def __init__(self, limit_max_requests=None, limit_max_requests_jitter=0):
pass
uvicorn_args: dict = {}
with patch("uvicorn.Config", _NewUvicornConfig):
ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter(
uvicorn_args=uvicorn_args,
max_requests_before_restart=1000,
jitter=50,
)
assert uvicorn_args["limit_max_requests_jitter"] == 50
def test_apply_uvicorn_jitter_skipped_on_old_uvicorn(self):
class _FakeUvicornConfig:
def __init__(self, limit_max_requests=None):
pass
uvicorn_args: dict = {}
with (
patch("uvicorn.Config", _FakeUvicornConfig),
patch("builtins.print") as mock_print,
):
ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter(
uvicorn_args=uvicorn_args,
max_requests_before_restart=1000,
jitter=50,
)
assert "limit_max_requests_jitter" not in uvicorn_args
assert any("0.41.0" in str(c) for c in mock_print.call_args_list)
def test_apply_uvicorn_jitter_without_base_warns(self):
uvicorn_args: dict = {}
with patch("builtins.print") as mock_print:
ProxyInitializationHelpers._apply_uvicorn_max_requests_jitter(
uvicorn_args=uvicorn_args,
max_requests_before_restart=None,
jitter=50,
)
assert "limit_max_requests_jitter" not in uvicorn_args
assert any("has no effect" in str(c) for c in mock_print.call_args_list)
@patch.dict(os.environ, {}, clear=True)
def test_construct_database_url_from_env_vars(self):
"""Test the construct_database_url_from_env_vars function with various scenarios"""
from litellm.proxy.utils import construct_database_url_from_env_vars
# Test with all required variables present
test_env = {
"DATABASE_HOST": "localhost:5432",
"DATABASE_USERNAME": "testuser",
"DATABASE_PASSWORD": "testpass",
"DATABASE_NAME": "testdb",
}
with patch.dict(os.environ, test_env):
result = construct_database_url_from_env_vars()
expected_url = "postgresql://testuser:testpass@localhost:5432/testdb"
assert result == expected_url
# Test with special characters that need URL encoding
test_env_special = {
"DATABASE_HOST": "localhost:5432",
"DATABASE_USERNAME": "user@with+special",
"DATABASE_PASSWORD": "test-password-special-chars",
"DATABASE_NAME": "db_name/test",
}
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"
assert result == expected_url
# Test without password (should still work)
test_env_no_password = {
"DATABASE_HOST": "localhost:5432",
"DATABASE_USERNAME": "testuser",
"DATABASE_NAME": "testdb",
}
with patch.dict(os.environ, test_env_no_password):
result = construct_database_url_from_env_vars()
expected_url = "postgresql://testuser@localhost:5432/testdb"
assert result == expected_url
# Test with missing required variables (should return None)
test_env_missing = {
"DATABASE_HOST": "localhost:5432",
"DATABASE_USERNAME": "testuser",
# Missing DATABASE_NAME
}
with patch.dict(os.environ, test_env_missing):
result = construct_database_url_from_env_vars()
assert result is None
# Test with empty environment (should return None)
with patch.dict(os.environ, {}, clear=True):
result = construct_database_url_from_env_vars()
assert result is None
@patch("uvicorn.run")
@patch("builtins.print")
def test_run_server_no_config_passed(self, mock_print, mock_uvicorn_run):
"""Test that run_server properly handles the case when no config is passed"""
import asyncio
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
mock_app = MagicMock()
mock_proxy_config = MagicMock()
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
# Mock the ProxyConfig.get_config method to return a proper async config
async def mock_get_config(config_file_path=None):
return {"general_settings": {}, "litellm_settings": {}}
mock_proxy_config_instance = MagicMock()
mock_proxy_config_instance.get_config = mock_get_config
mock_proxy_config.return_value = mock_proxy_config_instance
mock_proxy_server_module = MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
)
# Only remove DATABASE_URL and DIRECT_URL to prevent the database setup
# code path from running. Do NOT use clear=True as it removes PATH, HOME,
# etc., which causes imports inside run_server to break in CI (the real
# litellm.proxy.proxy_server import has heavy side effects that fail
# without a proper environment).
env_overrides = {
"DATABASE_URL": "",
"DIRECT_URL": "",
"IAM_TOKEN_DB_AUTH": "",
"AZURE_POSTGRESQL_AUTH": "",
"USE_AWS_KMS": "",
}
with patch.dict(os.environ, env_overrides):
# Remove DATABASE_URL entirely so the DB setup block is skipped
os.environ.pop("DATABASE_URL", None)
os.environ.pop("DIRECT_URL", None)
with (
patch.dict(
"sys.modules",
{"litellm.proxy.proxy_server": mock_proxy_server_module},
),
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",
"host": "localhost",
"port": 8000,
}
# Test with no config parameter (config=None)
result = runner.invoke(run_server, ["--local"])
assert result.exit_code == 0, (
f"run_server failed with exit_code={result.exit_code}, "
f"output={result.output}, exception={result.exception}"
)
# Verify that uvicorn.run was called
mock_uvicorn_run.assert_called_once()
# Reset mocks for second test
mock_uvicorn_run.reset_mock()
# Test with explicit --config None (should behave the same)
result = runner.invoke(run_server, ["--local", "--config", "None"])
assert result.exit_code == 0, (
f"run_server failed with exit_code={result.exit_code}, "
f"output={result.output}, exception={result.exception}"
)
# Verify that uvicorn.run was called again
mock_uvicorn_run.assert_called_once()
class TestQueryEngineReaperWiring:
def _invoke_run_server(self, args):
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")
}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
},
),
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,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = runner.invoke(run_server, args)
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"]
)
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"]
)
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()
@pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
def test_gunicorn_arbiter_starts_reaper(self):
pytest.importorskip("gunicorn")
with (
patch("gunicorn.app.base.BaseApplication.run"),
patch(
"litellm.proxy.proxy_cli.start_query_engine_reaper"
) as mock_start_reaper,
):
ProxyInitializationHelpers._run_gunicorn_server(
host="127.0.0.1",
port=4010,
app=MagicMock(),
num_workers=1,
ssl_certfile_path=None,
ssl_keyfile_path=None,
)
mock_start_reaper.assert_called_once()
class TestRunServerDbSetup:
"""Tests for run_server's prisma setup_database behavior."""
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
def test_use_prisma_db_push_flag_behavior(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""Test that use_prisma_db_push flag correctly controls PrismaManager.setup_database use_migrate parameter"""
from litellm.proxy.proxy_cli import run_server
# Mock subprocess.run to simulate prisma being available
mock_subprocess_run.return_value = MagicMock(returncode=0)
# Mock should_update_prisma_schema to return True (so setup_database gets called)
mock_should_update_schema.return_value = True
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
# Use standalone_mode=False to bypass Click's CliRunner stream
# isolation which causes flaky "I/O operation on closed file"
# errors in CI environments (Click 8.3.x stream lifecycle issue).
# Test 1: Without --use_prisma_db_push flag (default behavior)
# use_prisma_db_push should be False (default), so use_migrate should be True
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
mock_setup_database.assert_called_with(use_migrate=True, use_v2_resolver=True)
# Reset mocks
mock_setup_database.reset_mock()
mock_should_update_schema.reset_mock()
mock_should_update_schema.return_value = True
# Test 2: With --use_prisma_db_push flag set
# use_prisma_db_push should be True, so use_migrate should be False
run_server.main(
["--local", "--skip_server_startup", "--use_prisma_db_push"],
standalone_mode=False,
)
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
def test_migrations_run_when_the_prisma_cli_is_not_on_path(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
tmp_path,
capsys,
):
from litellm.proxy.proxy_cli import run_server
mock_should_update_schema.return_value = True
empty_bin = tmp_path / "emptybin"
empty_bin.mkdir()
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
clean_env["PATH"] = str(empty_bin)
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
patch( # test-quality-ok: same isolation as the sibling CLI tests above
"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",
"host": "localhost",
"port": 8000,
}
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
assert "prisma CLI is neither on PATH" not in capsys.readouterr().out
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
@pytest.mark.parametrize(
("database_url", "exits"),
(("postgresql://test:test@localhost:5432/test", True), (None, False)),
ids=("database-url-set", "no-database-url"),
)
@patch("atexit.register")
def test_startup_exits_when_the_prisma_toolchain_is_missing_only_if_a_database_is_configured(
self,
mock_atexit_register,
database_url,
exits,
tmp_path,
capsys,
):
"""A DATABASE_URL with no way to run the Prisma CLI is fatal; no DATABASE_URL needs no Prisma at all."""
from litellm_proxy_extras import prisma_toolchain
from litellm.proxy.proxy_cli import run_server
empty_bin = tmp_path / "emptybin"
empty_bin.mkdir()
real_find_spec = prisma_toolchain.importlib.util.find_spec
def hide_prisma(name, package=None):
return None if name == "prisma" else real_find_spec(name, package)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["PATH"] = str(empty_bin)
if database_url is not None:
clean_env["DATABASE_URL"] = database_url
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module},
),
patch.object(prisma_toolchain.importlib.util, "find_spec", side_effect=hide_prisma),
pytest.raises(SystemExit) if exits else nullcontext() as exit_info,
):
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
out = capsys.readouterr().out
if exits:
assert exit_info.value.code == 1
assert "a database URL is set but the prisma CLI is neither on PATH nor importable" in out
assert "pip install 'litellm[extra_proxy]'" in out
else:
assert "prisma CLI" not in out
assert "Setup complete" in out
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
def test_startup_fails_when_db_setup_fails(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""Test that proxy exits with code 1 when PrismaManager.setup_database returns False, with no opt-in flag"""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_should_update_schema.return_value = True
mock_setup_database.return_value = False
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
with pytest.raises(SystemExit) as exc_info:
run_server.main(["--local", "--skip_server_startup"], standalone_mode=False)
assert exc_info.value.code == 1
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
def test_startup_exits_on_non_postgres_database_url(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""A sqlite DATABASE_URL must exit immediately, before any prisma call,
instead of stalling on a migration against the postgresql-only schema."""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_should_update_schema.return_value = True
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "sqlite:///data/litellm.db"
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
):
with pytest.raises(SystemExit) as exc_info:
run_server.main(
["--local", "--skip_server_startup"], standalone_mode=False
)
assert exc_info.value.code == 1
mock_setup_database.assert_not_called()
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
def test_v2_migration_resolver_opts_in_via_env_var(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
capsys,
):
"""USE_V2_MIGRATION_RESOLVER=true must select the v2 resolver.
The Helm migrations Job runs `python litellm/proxy/prisma_migration.py`,
which calls run_server with a fixed argv, so a deployment reaches the
resolver through the env var rather than a CLI flag.
"""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_should_update_schema.return_value = True
mock_setup_database.return_value = True
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
clean_env["USE_V2_MIGRATION_RESOLVER"] = "true"
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
):
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
@pytest.mark.parametrize(
("arguments", "environment", "warned"),
(
(("--local", "--skip_server_startup", "--enforce_prisma_migration_check"), {}, True),
(("--local", "--skip_server_startup"), {"ENFORCE_PRISMA_MIGRATION_CHECK": "true"}, False),
(("--local", "--skip_server_startup"), {"ENFORCE_PRISMA_MIGRATION_CHECK": "false"}, False),
),
ids=("cli-flag", "env-true", "env-false"),
)
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True)
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes", return_value=True)
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=True)
def test_the_retired_enforce_prisma_migration_check_opt_in_still_parses_and_changes_nothing(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
arguments,
environment,
warned,
capsys,
):
"""Deployments still pass the flag or set the env var; the flag is accepted with a
deprecation line and the env var is ignored, and a successful setup boots either way."""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
with (
patch.dict(os.environ, {**clean_env, **environment}, clear=True),
patch.dict(
"sys.modules",
{"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module},
),
):
run_server.main(list(arguments), standalone_mode=False)
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
assert ("--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out) is warned
def test_the_retired_enforce_prisma_migration_check_opt_in_warns_on_health(
self, capsys: pytest.CaptureFixture[str]
) -> None:
"""`--health` returns before the config is read, where a server start emits the deprecation line;
it still prints it, as it did when the notice came before that return."""
with patch("httpx.get", return_value=httpx.Response(200, json={})) as mock_get:
run_server.main(["--health", "--enforce_prisma_migration_check"], standalone_mode=False)
mock_get.assert_called_once_with(url="http://0.0.0.0:4000/health")
assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out
def test_the_retired_enforce_prisma_migration_check_opt_in_warns_on_test(
self, capsys: pytest.CaptureFixture[str]
) -> None:
"""`--test` returns before the config is read too, and prints the deprecation line before its request."""
with pytest.raises(ValueError, match="Invalid test value"):
run_server.main(["--test", "--enforce_prisma_migration_check"], standalone_mode=False)
assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True)
def test_the_retired_enforce_prisma_migration_check_opt_in_warns_without_a_database(
self,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
capsys,
):
"""The deprecation line does not depend on reaching database setup: a deployment that
passes the flag with no DATABASE_URL still learns the flag is dead."""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module},
),
):
run_server.main(
["--local", "--skip_server_startup", "--enforce_prisma_migration_check"],
standalone_mode=False,
)
mock_setup_database.assert_not_called()
assert "--enforce_prisma_migration_check is deprecated and has no effect" in capsys.readouterr().out
@pytest.mark.parametrize(
"use_legacy_flag, env_value, expected",
[
(False, None, True),
(False, "true", True),
(False, "false", False),
(True, None, False),
(True, "true", False),
],
ids=[
"unset-env-defaults-to-v2",
"env-true-selects-v2",
"env-false-selects-v1",
"legacy-flag-selects-v1",
"legacy-flag-beats-env-true",
],
)
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
)
def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self):
from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli
assert deprecated_v2_flag_passed_on_cli() is False
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema")
def test_legacy_resolver_flag_reaches_database_setup(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""--use_legacy_migration_resolver must reach the database setup call.
The resolver decision itself is covered mock-free above; this is the
one wiring check that the flag is threaded through run_server.
"""
from litellm.proxy.proxy_cli import run_server
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_should_update_schema.return_value = True
mock_setup_database.return_value = True
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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", "USE_V2_MIGRATION_RESOLVER")
}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
):
run_server.main(
[
"--local",
"--skip_server_startup",
"--use_legacy_migration_resolver",
],
standalone_mode=False,
)
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=False)
@pytest.mark.parametrize(
("arguments", "migrated", "exits", "waits_for_the_build"),
(
(("--local", "--skip_server_startup"), True, True, True),
(("--local",), True, False, False),
(("--local",), False, True, False),
),
ids=("migration-job", "serving-proxy", "serving-proxy-whose-migrations-failed"),
)
@patch("uvicorn.run")
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database", return_value=True)
@patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes", return_value=False)
@patch("litellm.proxy.db.prisma_client.PrismaManager.start_request_log_index_build")
@patch("litellm.proxy.db.check_migration.check_prisma_schema_diff")
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=True)
def test_the_migration_job_waits_for_the_index_build_and_a_serving_proxy_starts_it_in_the_background(
self,
mock_should_update_schema,
mock_check_schema_diff,
mock_start_build,
mock_build_indexes,
mock_setup_database,
mock_atexit_register,
mock_subprocess_run,
mock_uvicorn_run,
arguments,
migrated,
exits,
waits_for_the_build,
):
"""`--skip_server_startup` is the migration job: it waits for the index build after the
migrations and exits 1 when one could not be built. A serving proxy that ran the
migrations starts the build in the background and serves whatever the build does; one
whose migrations failed exits 1 and starts no build."""
from litellm.proxy.proxy_cli import run_server
mock_setup_database.return_value = migrated
mock_subprocess_run.return_value = MagicMock(returncode=0)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
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["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
outcome = pytest.raises(SystemExit) if exits else nullcontext()
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"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,
outcome as exc_info,
):
mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000}
run_server.main(list(arguments), standalone_mode=False)
assert (exc_info is not None and exc_info.value.code == 1) is exits
mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True)
assert mock_build_indexes.call_count == int(migrated and waits_for_the_build)
assert mock_start_build.call_count == int(migrated and not waits_for_the_build)
# --- Module-level helpers for worker startup hook tests ---
_dummy_hook_called = False
def _dummy_hook():
"""A simple sync hook used by test_should_run_worker_startup_hooks."""
global _dummy_hook_called
_dummy_hook_called = True
_dummy_async_hook_called = False
async def _dummy_async_hook():
"""A simple async hook used by test_should_run_async_worker_startup_hook."""
global _dummy_async_hook_called
_dummy_async_hook_called = True
def _failing_hook():
"""A hook that always raises, used by test_should_raise_on_failing_hook."""
raise RuntimeError("Hook failed on purpose")
class TestWorkerStartupHooks:
"""Tests for the LITELLM_WORKER_STARTUP_HOOKS mechanism in proxy_startup_event."""
@pytest.mark.asyncio
async def test_should_run_worker_startup_hooks(self):
"""Sync worker startup hook is called during proxy_startup_event."""
global _dummy_hook_called
_dummy_hook_called = False
from litellm.proxy.proxy_server import proxy_startup_event
env_overrides = {
"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.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
try:
async with proxy_startup_event(app=None) as _:
pass
except Exception:
pass # We expect errors after the hook (no DB, etc.)
assert _dummy_hook_called is True, "Sync startup hook was not called"
@pytest.mark.asyncio
async def test_should_run_async_worker_startup_hook(self):
"""Async worker startup hook is awaited during proxy_startup_event."""
global _dummy_async_hook_called
_dummy_async_hook_called = False
from litellm.proxy.proxy_server import proxy_startup_event
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.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
try:
async with proxy_startup_event(app=None) as _:
pass
except Exception:
pass
assert _dummy_async_hook_called is True, "Async startup hook was not called"
@pytest.mark.asyncio
async def test_should_raise_on_failing_worker_startup_hook(self):
"""A failing worker startup hook propagates the error."""
from litellm.proxy.proxy_server import proxy_startup_event
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.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
with pytest.raises(RuntimeError, match="Hook failed on purpose"):
async with proxy_startup_event(app=None) as _:
pass
def test_should_skip_when_no_hooks_set(self):
"""When LITELLM_WORKER_STARTUP_HOOKS is not set, no hooks are executed."""
global _dummy_hook_called
_dummy_hook_called = False
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("LITELLM_WORKER_STARTUP_HOOKS", None)
# The hook block should be skipped entirely when env var is absent
assert "LITELLM_WORKER_STARTUP_HOOKS" not in os.environ
# Verify that an empty env var value also results in no hook execution
assert os.environ.get("LITELLM_WORKER_STARTUP_HOOKS", "") == ""
@pytest.mark.asyncio
async def test_should_run_multiple_hooks(self):
"""Multiple comma-separated hooks are all called."""
global _dummy_hook_called, _dummy_async_hook_called
_dummy_hook_called = False
_dummy_async_hook_called = False
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"
)
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.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
try:
async with proxy_startup_event(app=None) as _:
pass
except Exception:
pass
assert _dummy_hook_called is True, "First hook was not called"
assert _dummy_async_hook_called is True, "Second hook was not called"
@pytest.mark.xdist_group("proxy_cli")
class TestPostgresStatementTimeoutOptions:
"""A batch that outlives the Prisma client's HTTP read timeout keeps running
server side and holds its row locks until the database finishes it. Postgres
`statement_timeout` / `lock_timeout` are the only bound that ends that wait,
so they must survive the trip from general_settings into the connection URL.
"""
@pytest.mark.parametrize(
"existing, statement_timeout, lock_timeout, expected",
[
("", 60, 15, "-c statement_timeout=60000 -c lock_timeout=15000"),
("", 60, None, "-c statement_timeout=60000"),
("", None, 15, "-c lock_timeout=15000"),
("", None, None, ""),
("", 0.25, None, "-c statement_timeout=250"),
(
"-c search_path=app",
60,
15,
"-c search_path=app -c statement_timeout=60000 -c lock_timeout=15000",
),
(
"-c statement_timeout=5000",
60,
15,
"-c statement_timeout=5000 -c lock_timeout=15000",
),
(
"-cstatement_timeout=5000",
60,
15,
"-cstatement_timeout=5000 -c lock_timeout=15000",
),
(
"--statement_timeout=5000",
60,
15,
"--statement_timeout=5000 -c lock_timeout=15000",
),
(
"-c statement_timeout=5000",
60,
15,
"-c statement_timeout=5000 -c lock_timeout=15000",
),
],
ids=[
"both",
"statement_only",
"lock_only",
"neither",
"fractional_seconds",
"preserves_unrelated_option",
"pinned_spaced_wins",
"pinned_compact_wins",
"pinned_double_dash_wins",
"pinned_extra_spaces_wins",
],
)
def test_pg_options_with_timeouts(self, existing, statement_timeout, lock_timeout, expected):
from litellm.proxy.proxy_cli import _pg_options_with_timeouts
assert _pg_options_with_timeouts(existing, statement_timeout, lock_timeout) == expected
@pytest.mark.parametrize(
"existing, idle_timeout, expected",
[
("", 30, "-c statement_timeout=60000 -c lock_timeout=15000 -c idle_in_transaction_session_timeout=30000"),
("", None, "-c statement_timeout=60000 -c lock_timeout=15000"),
(
"-c idle_in_transaction_session_timeout=5000",
30,
"-c idle_in_transaction_session_timeout=5000 -c statement_timeout=60000 -c lock_timeout=15000",
),
],
ids=["idle_set", "idle_unset", "pinned_idle_wins"],
)
def test_pg_options_with_idle_in_transaction_timeout(
self,
existing: str,
idle_timeout: int | None,
expected: str,
) -> None:
"""A transaction that opened and then stalled holds its connection and its
locks for as long as the client stays silent; ``idle_in_transaction_session_timeout``
is the only server-side bound on that, so it rides the same ``options`` string."""
from litellm.proxy.proxy_cli import _pg_options_with_timeouts
assert _pg_options_with_timeouts(existing, 60, 15, idle_timeout) == expected
def test_timeouts_reach_the_database_url_from_general_settings(self, tmp_path):
"""The whole point of the setting: it has to land on DATABASE_URL."""
import yaml
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump(
{
"model_list": [],
"general_settings": {
"database_statement_timeout": 60,
"database_lock_timeout": 15,
"database_idle_in_transaction_session_timeout": 30,
},
}
)
)
modified_url = self._run_server_and_capture_database_url(str(config_path))
options = urlparse.parse_qs(urlparse.urlparse(modified_url).query)["options"][0]
assert "-c statement_timeout=60000" in options
assert "-c lock_timeout=15000" in options
assert "-c idle_in_transaction_session_timeout=30000" in options
def test_no_options_param_when_unset(self, tmp_path):
"""Unset must mean today's behavior, not an empty options string."""
import yaml
config_path = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": {}}))
modified_url = self._run_server_and_capture_database_url(str(config_path))
assert "options" not in urlparse.parse_qs(urlparse.urlparse(modified_url).query)
def test_direct_url_is_never_bounded(self, tmp_path):
"""DIRECT_URL serves migrations, which must not be cancelled mid-way."""
import yaml
config_path = tmp_path / "config.yaml"
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"
)
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)
def test_non_numeric_timeout_fails_fast(self, tmp_path):
"""A mistyped value must fail at startup, not deep inside URL assembly."""
import yaml
from pydantic import ValidationError
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": "sixty"}})
)
with pytest.raises(ValidationError):
self._run_server_and_capture_database_url(str(config_path))
def test_operator_pinned_url_options_are_preserved(self, tmp_path):
"""An operator's own `options` on DATABASE_URL must not be clobbered."""
import yaml
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump(
{
"model_list": [],
"general_settings": {"database_statement_timeout": 60},
}
)
)
modified_url = self._run_server_and_capture_database_url(
str(config_path),
database_url="postgresql://t:t@localhost:5432/t?options=-c%20search_path%3Dapp",
)
options = urlparse.parse_qs(urlparse.urlparse(modified_url).query)["options"][0]
assert "-c search_path=app" in options
assert "-c statement_timeout=60000" in options
@staticmethod
def _run_server_and_capture_database_url(
config_path: str,
database_url: str = "postgresql://t:t@localhost:5432/t",
) -> str:
return _run_server_and_capture_urls(config_path, database_url=database_url)["DATABASE_URL"]
_CAPTURED_DB_ENV_VARS = ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA")
def _run_server_and_capture_urls(
config_path: str,
database_url: str = "postgresql://t:t@localhost:5432/t",
direct_url: str | None = None,
read_replica_url: str | None = None,
extra_args: tuple[str, ...] = (),
) -> dict:
loaded_config = yaml.safe_load(Path(config_path).read_text())
mock_proxy_config = MagicMock()
mock_proxy_config.return_value.get_config = AsyncMock(return_value=loaded_config)
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=mock_proxy_config,
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in _CAPTURED_DB_ENV_VARS}
clean_env["DATABASE_URL"] = database_url
if direct_url is not None:
clean_env["DIRECT_URL"] = direct_url
if read_replica_url is not None:
clean_env["DATABASE_URL_READ_REPLICA"] = read_replica_url
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
patch("subprocess.run", return_value=MagicMock(returncode=0)),
patch("atexit.register"),
patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False),
patch("litellm.proxy.db.check_migration.check_prisma_schema_diff"),
):
run_server.main(
["--config", config_path, "--local", "--skip_server_startup", *extra_args],
standalone_mode=False,
)
return {k: os.environ[k] for k in _CAPTURED_DB_ENV_VARS if k in os.environ}
class TestReadReplicaConnectionParams:
"""The reader is a second Prisma client with its own pool. Without the
configured params on DATABASE_URL_READ_REPLICA it sizes itself from Prisma's
`num_physical_cpus * 2 + 1` default, so an operator's cap is not the cap that
gets enforced.
"""
def test_pool_settings_reach_the_read_replica_url(self, tmp_path):
import yaml
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump(
{
"model_list": [],
"general_settings": {
"database_connection_pool_limit": 3,
"database_connection_pool_timeout": 20,
"database_connect_timeout": 15,
"database_socket_timeout": 120,
"database_disable_prepared_statements": True,
"database_statement_timeout": 60,
},
}
)
)
captured = _run_server_and_capture_urls(
str(config_path),
read_replica_url="postgresql://t:t@reader:5432/t",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["connection_limit"] == ["3"]
assert query["pool_timeout"] == ["20"]
assert query["connect_timeout"] == ["15"]
assert query["socket_timeout"] == ["120"]
assert query["pgbouncer"] == ["true"]
assert "-c statement_timeout=60000" in query["options"][0]
def test_operator_pinned_replica_params_win(self, tmp_path):
"""The documented workaround (params pinned on the replica URL) must keep
working, so an operator who tuned the reader separately is not overridden.
"""
import yaml
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump(
{
"model_list": [],
"general_settings": {
"database_connection_pool_limit": 3,
"database_connection_pool_timeout": 20,
},
}
)
)
captured = _run_server_and_capture_urls(
str(config_path),
read_replica_url="postgresql://t:t@reader:5432/t?connection_limit=50",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["connection_limit"] == ["50"]
assert query["pool_timeout"] == ["20"]
def test_connection_budget_line_counts_the_limits_the_final_urls_carry(
self,
tmp_path: Path,
capsys: pytest.CaptureFixture[str],
) -> None:
import yaml
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})
)
_run_server_and_capture_urls(
str(config_path),
read_replica_url="postgresql://t:t@reader:5432/t?connection_limit=50",
)
assert (
"1 worker(s) x (writer connection_limit 3 + reader connection_limit 50) = up to 53 connections"
in capsys.readouterr().out
)
def test_connection_budget_line_counts_one_worker_under_hypercorn(
self,
tmp_path: Path,
capsys: pytest.CaptureFixture[str],
) -> None:
import yaml
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})
)
_run_server_and_capture_urls(
str(config_path),
extra_args=("--run_hypercorn", "--num_workers", "4"),
)
assert "1 worker(s) x writer connection_limit 3 = up to 3 connections" in capsys.readouterr().out
def test_extra_connection_params_never_carry_a_schema_override_to_the_reader(self, tmp_path):
"""database_extra_connection_params is an untyped passthrough, so it can carry a
search_path. The writer keeps it, the reader must not inherit it, or replica
queries resolve against the writer's schema.
"""
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.dump(
{
"model_list": [],
"general_settings": {
"database_connection_pool_limit": 3,
"database_extra_connection_params": {
"options": "-c search_path=writer_schema",
"schema": "writer_schema",
"socket_timeout": 90,
},
},
}
)
)
captured = _run_server_and_capture_urls(
str(config_path),
read_replica_url="postgresql://t:t@reader:5432/t",
)
writer_query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert writer_query["options"] == ["-c search_path=writer_schema"]
assert writer_query["schema"] == ["writer_schema"]
reader_query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert reader_query["connection_limit"] == ["3"]
assert reader_query["socket_timeout"] == ["90"]
assert "options" not in reader_query
assert "schema" not in reader_query
def test_replica_url_untouched_when_unset(self, tmp_path):
import yaml
config_path = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": {}}))
captured = _run_server_and_capture_urls(str(config_path))
assert "DATABASE_URL_READ_REPLICA" not in captured
class TestMaxIdleConnectionLifetimeDefault:
"""The proxy defaults `max_idle_connection_lifetime` below common infra idle
timeouts so stale pooled connections are recycled instead of failing requests."""
def _config(self, tmp_path, general_settings):
config_path = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": general_settings}))
return str(config_path)
def test_default_applied_to_database_and_direct_url(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
direct_url="postgresql://t:t@localhost:5432/t",
)
for env_var in ("DATABASE_URL", "DIRECT_URL"):
query = urlparse.parse_qs(urlparse.urlparse(captured[env_var]).query)
assert query["max_idle_connection_lifetime"] == ["60"], env_var
def test_url_pinned_value_wins_over_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
database_url="postgresql://t:t@localhost:5432/t?max_idle_connection_lifetime=300",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["300"]
def test_url_pinned_value_wins_over_config_key(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
database_url="postgresql://t:t@localhost:5432/t?max_idle_connection_lifetime=300",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["300"]
def test_config_key_overrides_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["45"]
def test_extra_connection_params_override_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(
tmp_path,
{"database_extra_connection_params": {"max_idle_connection_lifetime": 120}},
),
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert query["max_idle_connection_lifetime"] == ["120"]
def test_read_replica_gets_the_default(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
read_replica_url="postgresql://t:t@reader:5432/t",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["max_idle_connection_lifetime"] == ["60"]
def test_replica_pinned_value_wins(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
read_replica_url="postgresql://t:t@reader:5432/t?max_idle_connection_lifetime=200",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["max_idle_connection_lifetime"] == ["200"]
def test_config_key_reaches_the_read_replica(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(tmp_path, {"database_max_idle_connection_lifetime": 45}),
read_replica_url="postgresql://t:t@reader:5432/t",
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL_READ_REPLICA"]).query)
assert query["max_idle_connection_lifetime"] == ["45"]
def test_idle_lifetime_params_prefers_configured_value(self):
from litellm.proxy.db.db_url_settings import idle_lifetime_params
assert dict(idle_lifetime_params(45)) == {"max_idle_connection_lifetime": 45}
assert dict(idle_lifetime_params(None)) == {"max_idle_connection_lifetime": 60}
class TestTokenAuthCliFlags:
"""`--azure_postgresql_auth` has to reach the URL assembly the same way the env var does."""
def _invoke_with_azure_host(self, args):
from click.testing import CliRunner
from litellm.proxy.db.token_auth import build_azure_entra_token_provider
from litellm.proxy.proxy_cli import run_server
build_azure_entra_token_provider.cache_clear()
clean_env = {
k: v
for k, v in os.environ.items()
if k
not in (
"DATABASE_URL",
"DIRECT_URL",
"IAM_TOKEN_DB_AUTH",
"AZURE_POSTGRESQL_AUTH",
"DATABASE_URL_READ_REPLICA",
)
}
clean_env["DATABASE_HOST"] = "writer.postgres.database.azure.com"
clean_env["DATABASE_USER"] = "litellm@contoso.onmicrosoft.com"
clean_env["DATABASE_NAME"] = "litellm_db"
mock_proxy_module = MagicMock(
app=MagicMock(),
ProxyConfig=MagicMock(),
KeyManagementSettings=MagicMock(),
save_worker_config=MagicMock(),
)
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
"sys.modules",
{
"proxy_server": mock_proxy_module,
"litellm.proxy.proxy_server": mock_proxy_module,
},
),
patch(
"litellm.secret_managers.get_azure_ad_token_provider.get_azure_ad_token_provider",
return_value=lambda: "ENTRA_TOKEN",
),
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,
):
mock_get_args.return_value = {
"app": "litellm.proxy.proxy_server:app",
"host": "localhost",
"port": 8000,
}
result = CliRunner().invoke(run_server, args)
database_url = os.getenv("DATABASE_URL")
toggle = os.getenv("AZURE_POSTGRESQL_AUTH")
build_azure_entra_token_provider.cache_clear()
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"]
)
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
assert database_url is not None
assert "ENTRA_TOKEN" in database_url
assert "writer.postgres.database.azure.com" in database_url
assert toggle == "True"
def test_without_the_flag_no_token_is_minted(self):
result, database_url, toggle = self._invoke_with_azure_host(["--local"])
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
assert "ENTRA_TOKEN" not in (database_url or "")
assert toggle is None
class TestLibpqSslParamTranslation:
"""Prisma ignores ``sslrootcert`` and treats ``sslmode=verify-full`` as
``prefer``, so the URL handed to it must carry Prisma's own strict dialect."""
@staticmethod
def _config(tmp_path, general_settings):
config_path = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": general_settings}))
return str(config_path)
def test_libpq_url_is_translated_on_every_prisma_url(self, tmp_path):
libpq = "?sslmode=verify-full&sslrootcert=/certs/rds-bundle.pem"
captured = _run_server_and_capture_urls(
self._config(tmp_path, {}),
database_url=f"postgresql://t:t@localhost:5432/t{libpq}",
direct_url=f"postgresql://t:t@direct:5432/t{libpq}",
read_replica_url=f"postgresql://t:t@reader:5432/t{libpq}",
)
for env_var in ("DATABASE_URL", "DIRECT_URL", "DATABASE_URL_READ_REPLICA"):
query = urlparse.parse_qs(urlparse.urlparse(captured[env_var]).query)
assert "sslrootcert" not in query, env_var
assert query["sslmode"] == ["require"], env_var
assert query["sslcert"] == ["/certs/rds-bundle.pem"], env_var
assert query["sslaccept"] == ["strict"], env_var
def test_libpq_params_from_extra_connection_params_are_translated(self, tmp_path):
captured = _run_server_and_capture_urls(
self._config(
tmp_path,
{
"database_extra_connection_params": {
"sslmode": "verify-full",
"sslrootcert": "/certs/rds-bundle.pem",
}
},
),
)
query = urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert "sslrootcert" not in query
assert query["sslmode"] == ["require"]
assert query["sslcert"] == ["/certs/rds-bundle.pem"]
assert query["sslaccept"] == ["strict"]
@pytest.mark.xdist_group("proxy_cli")
class TestValidateConfigFlag:
def test_validate_config_valid_config_exits_zero(self, tmp_path, monkeypatch):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
monkeypatch.delenv("DATABASE_URL", raising=False)
monkeypatch.delenv("DIRECT_URL", raising=False)
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.safe_dump(
{
"model_list": [
{
"model_name": "gpt-4o",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "sk-fake",
},
}
]
}
)
)
result = CliRunner().invoke(run_server, ["--config", str(config_path), "--validate_config"])
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
assert "config OK" in result.output
def test_validate_config_invalid_mcp_server_exits_one(self, tmp_path, monkeypatch):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
monkeypatch.delenv("DATABASE_URL", raising=False)
monkeypatch.delenv("DIRECT_URL", raising=False)
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.safe_dump(
{
"mcp_servers": {
"zapier": {
"url": "https://example.com/mcp",
"transport": "http",
"per_server_oauth_discovery": "yes",
}
}
}
)
)
result = CliRunner().invoke(run_server, ["--config", str(config_path), "--validate_config"])
assert result.exit_code == 1, f"exit_code={result.exit_code}, output={result.output}"
assert "per_server_oauth_discovery must be a boolean" in result.output
def test_validate_config_without_config_is_usage_error(self, monkeypatch):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
result = CliRunner().invoke(run_server, ["--validate_config"])
assert result.exit_code != 0
assert "--validate_config requires --config" in result.output
@patch("subprocess.Popen")
def test_validate_config_with_ollama_model_does_not_start_ollama(self, mock_popen, tmp_path, monkeypatch):
from click.testing import CliRunner
from litellm.proxy.proxy_cli import run_server
monkeypatch.delenv("DATABASE_URL", raising=False)
monkeypatch.delenv("DIRECT_URL", raising=False)
config_path = tmp_path / "config.yaml"
config_path.write_text(
yaml.safe_dump(
{
"model_list": [
{
"model_name": "gpt-4o",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "sk-fake",
},
}
]
}
)
)
result = CliRunner().invoke(
run_server,
["--config", str(config_path), "--model", "ollama/llama3", "--validate_config"],
)
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
assert "config OK" in result.output
mock_popen.assert_not_called()
def _boot_from_config_file(config_path: Path, *flags: str) -> str | None:
env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "JSON_LOGS")}
with patch.dict(os.environ, env, clear=True):
run_server.main(["--config", str(config_path), "--skip_server_startup", *flags], standalone_mode=False)
return os.environ.get("JSON_LOGS")
class TestJsonLogsConfigReachesSpawnedWorkers:
"""`litellm_settings.json_logs: true` turns JSON on inside the CLI process, but the workers that
uvicorn or granian spawn are fresh interpreters that pick plain or JSON from the JSON_LOGS
environment variable at import. The CLI has to export it, or every worker logs plain text until
its own load_config reaches the key."""
def test_json_logs_true_exports_json_logs_to_spawned_workers(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "json_logs", False)
monkeypatch.setattr(sys, "excepthook", sys.excepthook)
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"json_logs": True}}))
assert _boot_from_config_file(config_path) == "true"
def test_plain_logs_config_leaves_json_logs_unset(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "json_logs", False)
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {}}))
assert _boot_from_config_file(config_path) is None
class TestEnforcePrismaMigrationCheckNoticeAroundTheConfigRead:
"""With a config file, the `--enforce_prisma_migration_check` deprecation notice waits until the config
is read, so a `json_logs: true` set only there makes it a JSON record. It is printed once either way,
and a boot whose config read fails still prints it before the error, as it did before the wait."""
NOTICE: Final = "--enforce_prisma_migration_check is deprecated and has no effect"
FLAG: Final = "--enforce_prisma_migration_check"
def test_a_config_read_that_fails_still_prints_the_notice(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.setattr(litellm, "json_logs", False)
with pytest.raises(Exception, match="Config file not found"):
_boot_from_config_file(tmp_path / "missing.yaml", self.FLAG)
assert capsys.readouterr().out.count(self.NOTICE) == 1
def test_a_plain_config_prints_the_notice_once(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.setattr(litellm, "json_logs", False)
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {}}))
_boot_from_config_file(config_path, self.FLAG)
assert capsys.readouterr().out.count(self.NOTICE) == 1
def test_a_config_only_json_logs_makes_the_notice_one_json_record(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
monkeypatch.setattr(litellm, "json_logs", False)
monkeypatch.setattr(sys, "excepthook", sys.excepthook)
config_path: Final = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "litellm_settings": {"json_logs": True}}))
_boot_from_config_file(config_path, self.FLAG)
notice_lines: Final = [line for line in capsys.readouterr().out.splitlines() if self.NOTICE in line]
assert len(notice_lines) == 1, notice_lines
assert self.NOTICE in json.loads(notice_lines[0])["message"]
def test_granian_flag_notices_are_json_records_in_json_mode(
monkeypatch: pytest.MonkeyPatch, capfd: pytest.CaptureFixture[str]
) -> None:
"""The yellow notices for the flags Granian ignores (`--max_requests_before_restart`, `--ciphers`) are
printed while the launcher starts; with json_logs on they are INFO records on stdout like every other
line, not plaintext between the records."""
pytest.importorskip("granian")
monkeypatch.setattr(litellm, "json_logs", True)
monkeypatch.setattr(sys, "excepthook", sys.excepthook)
turn_on_json()
with (
patch("granian.Granian") as mock_granian_cls,
patch("granian.constants.Interfaces", SimpleNamespace(ASGI="asgi")),
):
ProxyInitializationHelpers._init_granian_server(
host="0.0.0.0",
port=4000,
num_workers=2,
ssl_certfile_path=None,
ssl_keyfile_path=None,
max_requests_before_restart=100,
ciphers="ECDHE",
granian_runtime_threads=None,
)
mock_granian_cls.return_value.serve.assert_called_once()
out, err = capfd.readouterr()
assert err == ""
records: Final = [json.loads(raw) for raw in out.splitlines() if raw.strip()]
assert [record["message"] for record in records] == [
"LiteLLM Proxy: Starting server on 0.0.0.0:4000 using Granian",
"LiteLLM: --max_requests_before_restart is not supported by Granian "
"(Granian uses workers_lifetime in seconds, not a per-request limit).",
"LiteLLM: --ciphers is not applied when using --run_granian.",
]
assert {record["level"] for record in records} == {"INFO"}
Reference in a new issue View git blame Copy permalink