From aae2deb839c8f39f82a58140f2ffa97b42d97a61 Mon Sep 17 00:00:00 2001
From: Harshit28j <Harshitjain0562@gmail.com>
Date: Tue, 10 Mar 2026 18:15:11 +0530
Subject: [PATCH] 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>
---
 litellm/proxy/proxy_cli.py                 |  2 -
 tests/test_litellm/proxy/test_proxy_cli.py | 57 ++++++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py
index 5f6ccfd9ba..be2f5ac7c1 100644
--- a/litellm/proxy/proxy_cli.py
+++ b/litellm/proxy/proxy_cli.py
@@ -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"
diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py
index cf6511c18a..bed77b43c1 100644
--- a/tests/test_litellm/proxy/test_proxy_cli.py
+++ b/tests/test_litellm/proxy/test_proxy_cli.py
@@ -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 ---