mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 10:21:52 +00:00
fix: remove redundant import and add test for startup failure
- Remove redundant `import sys` (already imported at module level) - Add test_startup_fails_when_db_setup_fails verifying sys.exit(1) when PrismaManager.setup_database returns False Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
4714128da5
commit
aae2deb839
@@ -854,8 +854,6 @@ def run_server( # noqa: PLR0915
|
||||
check_prisma_schema_diff(db_url=None)
|
||||
else:
|
||||
if not PrismaManager.setup_database(use_migrate=not use_prisma_db_push):
|
||||
import sys
|
||||
|
||||
print( # noqa
|
||||
"\033[1;31mLiteLLM Proxy: Database setup failed after multiple retries. "
|
||||
"The proxy cannot start safely. Please check your database connection and migration status.\033[0m"
|
||||
|
||||
@@ -664,6 +664,63 @@ class TestHealthAppFactory:
|
||||
)
|
||||
mock_setup_database.assert_called_with(use_migrate=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"""
|
||||
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
|
||||
|
||||
|
||||
# --- Module-level helpers for worker startup hook tests ---
|
||||
|
||||
|
||||
Reference in New Issue
Block a user