From 1b8f704035a358d31962ca24d8ceda77d3dde935 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:17:18 -0700 Subject: [PATCH] fix(proxy): await the cancelled view setup task quietly and assert it starts at boot Use contextlib.suppress for the cancelled task in stop_view_setup_task, make the legacy prisma setup test inject a plain mock for the synchronous start_view_setup_task and assert it is called, and drop the docstrings the branch added to tests --- litellm/proxy/utils.py | 4 +--- .../spend_tracking/spend_e2e_client.py | 2 -- tests/proxy_unit_tests/test_proxy_server.py | 17 +++++------------ tests/test_litellm/proxy/test_proxy_server.py | 7 ------- .../test_prisma_client_lifecycle.py | 8 -------- 5 files changed, 6 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 7f5609b2e39..37c5c8acec8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6350,10 +6350,8 @@ class PrismaClient: if self._view_setup_task is None: return self._view_setup_task.cancel() - try: + with contextlib.suppress(asyncio.CancelledError): await self._view_setup_task - except asyncio.CancelledError: - pass self._view_setup_task = None async def _run_view_setup( diff --git a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py index 233aa81c2af..b7f59fe5f89 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -360,8 +360,6 @@ class SpendClient: return self.proxy.transport.probe(path, params=params) def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult: - """Re-probe a route that depends on startup work the proxy finishes after it - starts serving, such as the spend views it creates once migrations land.""" outcome: Final = await_converged( lambda: self.probe(path, params=params), converged=lambda result: result.healthy, diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 47792b90b08..ed0380058a5 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2382,7 +2382,7 @@ async def test_proxy_model_group_info_rerank(prisma_client): # noqa: F811 # py @pytest.mark.asyncio async def test_proxy_server_prisma_setup(): - from litellm.proxy.proxy_server import ProxyStartupEvent, proxy_state + from litellm.proxy.proxy_server import ProxyStartupEvent from litellm.proxy.utils import ProxyLogging from litellm.caching import DualCache @@ -2393,35 +2393,28 @@ async def test_proxy_server_prisma_setup(): ) as mock_prisma_client: mock_client = mock_prisma_client.return_value # This is the mocked instance mock_client.connect = AsyncMock() # Mock the connect method - mock_client.check_view_exists = AsyncMock() # Mock the check_view_exists method + mock_client.start_view_setup_task = MagicMock() mock_client.health_check = AsyncMock() # Mock the health_check method - mock_client._set_spend_logs_row_count_in_proxy_state = ( - AsyncMock() - ) # Mock the _set_spend_logs_row_count_in_proxy_state method mock_client.start_db_health_watchdog_task = AsyncMock() # Mock the db attribute with start_token_refresh_task for RDS IAM token refresh mock_db = MagicMock() mock_db.start_token_refresh_task = AsyncMock() mock_client.db = mock_db - await ProxyStartupEvent._setup_prisma_client( + prisma_client = await ProxyStartupEvent._setup_prisma_client( database_url=os.getenv("DATABASE_URL"), proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), user_api_key_cache=user_api_key_cache, ) - # Verify our mocked methods were called + assert prisma_client is mock_client mock_client.connect.assert_called_once() - mock_client.check_view_exists.assert_called_once() + mock_client.start_view_setup_task.assert_called_once() # Note: This is REALLY IMPORTANT to check that the health check is called # This is how we ensure the DB is ready before proceeding mock_client.health_check.assert_called_once() - # check that the spend logs row count is set in proxy state - mock_client._set_spend_logs_row_count_in_proxy_state.assert_called_once() - assert proxy_state.get_proxy_state_variable("spend_logs_row_count") is not None - @pytest.mark.asyncio async def test_proxy_server_prisma_setup_invalid_db(monkeypatch): diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7ab142c5fd0..67f386ee457 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1636,9 +1636,6 @@ class _ShutdownAwarePrisma(MockPrisma): @pytest.mark.asyncio async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): - """The view setup task keeps polling for the spend-log table while migrations - run, so a shutdown inside that window has to cancel it rather than leave it - to die with the event loop.""" import yaml from fastapi import FastAPI @@ -13415,10 +13412,6 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch @pytest.mark.asyncio async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch): - """View creation used to be two fire-and-forget ``asyncio.create_task`` calls - that raised and vanished when the migrations Job had not created - ``LiteLLM_SpendLogs`` yet (LIT-5211). Startup must hand the work to the client's - held task, which waits for the table, and must not call the two coroutines directly.""" monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") mock_client = _mock_startup_prisma_client() diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index bb34486771c..c31713d5802 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -236,9 +236,6 @@ async def test_disconnect_raises_when_underlying_fails( async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views( prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch ) -> None: - """On a fresh database the migrations Job can still be running when the proxy - boots. The views reference ``LiteLLM_SpendLogs``, so creating them before the - table exists raised inside a fire-and-forget task and the views never appeared.""" monkeypatch.delenv("DATABASE_SCHEMA", raising=False) probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) call_order = _wire_view_setup(prisma_client, probe) @@ -293,9 +290,6 @@ async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: P @pytest.mark.asyncio async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None: - """``LiteLLM_SpendLogs`` lands early in the migration set while - ``LiteLLM_VerificationTokenView`` references columns the newest migrations add, - so the first attempt after the table appears can still fail.""" probe = AsyncMock(return_value=_present()) call_order = _wire_view_setup(prisma_client, probe) prisma_client.check_view_exists.side_effect = [RuntimeError('column "tpd_limit" does not exist'), None] @@ -358,8 +352,6 @@ async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( async def test_run_view_setup_reports_the_last_error_when_views_keep_failing_on_a_present_table( prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture ) -> None: - """A database role without CREATE on the schema fails every attempt even though - the table is there, so the timeout must blame that error, not missing migrations.""" _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public")