From 8b5b3438417902f8b71226ade63c6b7ec2dd113a Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 23 Jan 2026 12:10:08 -0800
Subject: [PATCH] attempt fix flaky tests

---
 tests/local_testing/test_router_utils.py   |  2 +-
 tests/test_litellm/proxy/test_proxy_cli.py | 26 ++++++++--------------
 2 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/tests/local_testing/test_router_utils.py b/tests/local_testing/test_router_utils.py
index 0e3835a7f9..55b1893c07 100644
--- a/tests/local_testing/test_router_utils.py
+++ b/tests/local_testing/test_router_utils.py
@@ -288,7 +288,7 @@ async def test_call_router_callbacks_on_failure():
                 mock_response="litellm.RateLimitError",
                 num_retries=0,
             )
-        await asyncio.sleep(1)
+        await asyncio.sleep(3)
         print(mock_callback.call_args_list)
         assert mock_callback.call_count == 1
 
diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py
index 5f03ef1817..f52f14b860 100644
--- a/tests/test_litellm/proxy/test_proxy_cli.py
+++ b/tests/test_litellm/proxy/test_proxy_cli.py
@@ -215,28 +215,21 @@ class TestProxyInitializationHelpers:
             assert "pool_timeout=60" in modified_url
 
     @patch("uvicorn.run")
-    @patch("builtins.print")
-    def test_skip_server_startup(self, mock_print, mock_uvicorn_run):
-        """Test that the skip_server_startup flag prevents server startup when True"""
+    @patch("atexit.register")  # 🔥 critical
+    def test_skip_server_startup(self, mock_atexit_register, mock_uvicorn_run):
         from click.testing import CliRunner
-
         from litellm.proxy.proxy_cli import run_server
 
         runner = CliRunner()
 
-        mock_app = MagicMock()
-        mock_proxy_config = MagicMock()
-        mock_key_mgmt = MagicMock()
-        mock_save_worker_config = MagicMock()
-
         with patch.dict(
             "sys.modules",
             {
                 "proxy_server": MagicMock(
-                    app=mock_app,
-                    ProxyConfig=mock_proxy_config,
-                    KeyManagementSettings=mock_key_mgmt,
-                    save_worker_config=mock_save_worker_config,
+                    app=MagicMock(),
+                    ProxyConfig=MagicMock(),
+                    KeyManagementSettings=MagicMock(),
+                    save_worker_config=MagicMock(),
                 )
             },
         ), patch(
@@ -248,16 +241,15 @@ class TestProxyInitializationHelpers:
                 "port": 8000,
             }
 
+            # --- skip startup ---
             result = runner.invoke(run_server, ["--local", "--skip_server_startup"])
 
             assert result.exit_code == 0
+            assert "Skipping server startup" in result.output
             mock_uvicorn_run.assert_not_called()
-            mock_print.assert_any_call(
-                "LiteLLM: Setup complete. Skipping server startup as requested."
-            )
 
+            # --- normal startup ---
             mock_uvicorn_run.reset_mock()
-            mock_print.reset_mock()
 
             result = runner.invoke(run_server, ["--local"])