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

@@ -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(

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,

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):

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()

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")