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:
Harshit28j
2026-03-10 18:15:11 +05:30
co-authored by Claude Opus 4.6
parent 4714128da5
commit aae2deb839
2 changed files with 57 additions and 2 deletions
-2
View File
@@ -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 ---