mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Fix: Currently database_connection_timeout pulls value from the database_connection_pool_timeout param from the general settings which does not adhere to the settings in the doc. Add a fix to consider either database_connection_timeout or database_connection_pool_timeout
This commit is contained in:
parent
5f49f29f4e
commit
2ea0e53ec8
2 changed files with 81 additions and 3 deletions
|
|
@ -803,8 +803,11 @@ def run_server( # noqa: PLR0915
|
|||
LiteLLMDatabaseConnectionPool.database_connection_pool_limit.value,
|
||||
)
|
||||
db_connection_timeout = general_settings.get(
|
||||
"database_connection_pool_timeout",
|
||||
LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value,
|
||||
"database_connection_timeout",
|
||||
general_settings.get(
|
||||
"database_connection_pool_timeout",
|
||||
LiteLLMDatabaseConnectionPool.database_connection_pool_timeout.value,
|
||||
),
|
||||
)
|
||||
if database_url and database_url.startswith("os.environ/"):
|
||||
original_dir = os.getcwd()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import fastapi
|
||||
import pytest
|
||||
|
|
@ -420,6 +420,81 @@ class TestProxyInitializationHelpers:
|
|||
call_args = mock_uvicorn_run.call_args
|
||||
assert call_args[1]["limit_max_requests"] == 123
|
||||
|
||||
@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_database_connection_timeout_key_is_used_for_pool_timeout(
|
||||
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_connection_pool_limit": 5,
|
||||
"database_connection_timeout": 30,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
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"] == 30
|
||||
|
||||
@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"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue