From 953a8b1b8552dcd72a97a669f9985156b4206bab Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 15 May 2026 16:51:45 -0700 Subject: [PATCH] test(proxy): isolate run_server CLI tests from prisma DB-setup path test_keepalive_timeout_flag and test_timeout_worker_healthcheck_flag were the only run_server tests in test_proxy_cli.py that neither stripped DATABASE_URL/DIRECT_URL nor mocked the prisma DB path. When a DATABASE_URL is present (CI/env leak), run_server --local enters the DB block and blocks in the un-timeout'd subprocess.run(["prisma"]) at proxy_cli.py:987 plus the ProxyExtrasDBManager migrate-deploy retry loops, ~370s per test on the CI runner. --dist=loadscope pins both to one xdist worker, so the proxy-infra job appears stuck at 99% and hits the 20-min timeout. Apply the same isolation every other run_server test in this file already uses: mock PrismaManager.setup_database + should_update_prisma_schema and strip DATABASE_URL/DIRECT_URL. Full module drops from 31.7s to 2.9s locally; both tests fall off the slow list. --- tests/test_litellm/proxy/test_proxy_cli.py | 38 ++++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 327200a6a9..46a55fc746 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -538,7 +538,13 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") - def test_keepalive_timeout_flag(self, mock_print, mock_uvicorn_run): + @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 @@ -551,7 +557,18 @@ class TestProxyInitializationHelpers: 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", { @@ -596,7 +613,13 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") - def test_timeout_worker_healthcheck_flag(self, mock_print, mock_uvicorn_run): + @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 @@ -609,7 +632,18 @@ class TestProxyInitializationHelpers: 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", {