scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 35
litellm/tests/test_litellm/proxy/test_proxy_cli.py
yuneng-jiang 6a0d03914c
test: drop the cwd-relative sys.path.insert calls from the test suite (#37802) 
* test: drop the cwd-relative sys.path.insert calls from the test suite

TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.

Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.

Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.

* test: drop the duplicate imports the sys.path sweep exposed to F811

* test(pre-call-utils): restore the os import the new bedrock tests need
2026-08-22 09:25:58 -07:00

2534 lines
93 KiB
Python
Raw Permalink Blame History

import inspect
import os
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import click
import fastapi
import pytest
import builtins
import types
import urllib.parse as urlparse
import uvicorn
import yaml
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
@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
# 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
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):
import litellm
monkeypatch.setenv("LITELLM_DEV_ENV_HOT_RELOAD", "True")
assert litellm._dev_env_hot_reload_enabled() is True
monkeypatch.setenv("LITELLM_DEV_ENV_HOT_RELOAD", "false")
assert litellm._dev_env_hot_reload_enabled() is False
monkeypatch.delenv("LITELLM_DEV_ENV_HOT_RELOAD", raising=False)
assert litellm._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"
@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
):
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")
}
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()
# --- 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()
@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()
@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
@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.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):
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_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,
):
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)
# 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 at line 820 of proxy_cli.py 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",
{
"proxy_server": MagicMock(
app=mock_app,
ProxyConfig=mock_proxy_config,
KeyManagementSettings=mock_key_mgmt,
save_worker_config=mock_save_worker_config,
),
# Also mock litellm.proxy.proxy_server to prevent the real
# import at line 820 of proxy_cli.py which has heavy side
# effects (FastAPI app init, logging setup, etc.)
"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.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_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=False
)
# 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=False
)
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@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_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""Test that proxy exits with code 1 when PrismaManager.setup_database returns False and --enforce_prisma_migration_check is set"""
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",
"--enforce_prisma_migration_check",
],
standalone_mode=False,
)
assert exc_info.value.code == 1
mock_setup_database.assert_called_once_with(
use_migrate=True, use_v2_resolver=False
)
@patch("subprocess.run")
@patch("atexit.register")
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
@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_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.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_setup_database,
mock_atexit_register,
mock_subprocess_run,
):
"""USE_V2_MIGRATION_RESOLVER 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 has no way to
pass --use_v2_migration_resolver and an env var is the only route in.
"""
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
)
# --- 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.test_litellm.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.test_litellm.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.test_litellm.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.test_litellm.proxy.test_proxy_cli:_dummy_hook,"
"tests.test_litellm.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
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,
},
}
)
)
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
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,
) -> 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"],
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_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 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
Reference in a new issue View git blame Copy permalink