diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e49d563328f..dae88da5974 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1510,6 +1510,12 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: except Exception as e: verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) + if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"): + try: + await prisma_client.stop_view_setup_task() + except Exception as e: + verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) + await _drain_spend_event_producer_on_shutdown() await flush_spend_counters_on_shutdown() @@ -10808,14 +10814,7 @@ class ProxyStartupEvent: if hasattr(prisma_client, "db") and hasattr(prisma_client.db, "start_token_refresh_task"): await prisma_client.db.start_token_refresh_task() - ## Add necessary views to proxy ## - asyncio.create_task( - prisma_client.check_view_exists() - ) # check if all necessary views exist. Don't block execution - - asyncio.create_task( - prisma_client._set_spend_logs_row_count_in_proxy_state() - ) # set the spend logs row count in proxy state. Don't block execution + prisma_client.start_view_setup_task() if hasattr(prisma_client, "start_db_health_watchdog_task"): await prisma_client.start_db_health_watchdog_task() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f6f437bea75..9de2b5fd282 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -38,6 +38,7 @@ from typing import ( Literal, Optional, Protocol, + TypeAlias, TypeVar, Union, cast, @@ -269,6 +270,15 @@ class _RelTuplesRow(TypedDict): reltuples: ReadOnly[int] +_VIEW_SETUP_POLL_INTERVAL_SECONDS: Final = 5.0 +_VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 +_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' +_VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) + +_ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] +_ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception + + class _EndUserBatchTable(Protocol): def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ... @@ -4297,6 +4307,7 @@ class PrismaClient: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() self._db_health_watchdog_task: asyncio.Task | None = None + self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None self._db_last_reconnect_attempt_ts: float = 0.0 self._db_reconnect_cooldown_seconds: int = max(1, int(os.getenv("PRISMA_RECONNECT_COOLDOWN_SECONDS", "15"))) self._db_read_only_recreate_ts: float = 0.0 @@ -6343,6 +6354,71 @@ class PrismaClient: self._db_health_watchdog_task = None verbose_proxy_logger.info("Stopped Prisma DB health watchdog") + def start_view_setup_task(self) -> None: + if self._view_setup_task is not None: + return + self._view_setup_task = asyncio.create_task(self._run_view_setup()) + + async def stop_view_setup_task(self) -> None: + if self._view_setup_task is None: + return + self._view_setup_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._view_setup_task + self._view_setup_task = None + + async def _run_view_setup( + self, + poll_interval_seconds: float = _VIEW_SETUP_POLL_INTERVAL_SECONDS, + deadline_seconds: float = _VIEW_SETUP_DEADLINE_SECONDS, + ) -> _ViewSetupOutcome: + deadline: Final = time.monotonic() + deadline_seconds + while True: + if (attempt := await self._attempt_view_setup()) == "ready": + return "ready" + if time.monotonic() >= deadline: + self._log_view_setup_timeout(attempt, deadline_seconds) + return "timed_out" + await asyncio.sleep(poll_interval_seconds) + + async def _attempt_view_setup(self) -> _ViewSetupAttempt: + try: + if not await self._view_setup_gate_table_present(): + verbose_proxy_logger.debug( + "Waiting for table %s before creating the spend views", _VIEW_SETUP_GATE_TABLE + ) + return "table_missing" + await self._set_spend_logs_row_count_in_proxy_state() + await self.check_view_exists() + return "ready" + except Exception as e: + verbose_proxy_logger.warning("Spend view setup attempt failed, retrying until the schema settles: %s", e) + return e + + def _log_view_setup_timeout( + self, last_attempt: Literal["table_missing"] | Exception, deadline_seconds: float + ) -> None: + if isinstance(last_attempt, Exception): + verbose_proxy_logger.error( + "Gave up creating the spend views after %ss; the last attempt failed with: %s. " + "Fix that error and restart the proxy.", + deadline_seconds, + last_attempt, + ) + return + verbose_proxy_logger.error( + "Gave up creating the spend views: table %s did not appear within %ss. " + "Run the database migrations against this database and restart the proxy.", + _VIEW_SETUP_GATE_TABLE, + deadline_seconds, + ) + + async def _view_setup_gate_table_present(self) -> bool: + rows: Final = _VIEW_SETUP_GATE_PROBE_ROWS.validate_python( + await self.db.query_raw("SELECT to_regclass($1) IS NOT NULL AS present", _VIEW_SETUP_GATE_TABLE) + ) + return rows[0]["present"] + async def _db_health_watchdog_loop(self) -> None: while True: try: 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 9ac97f57f47..b7f59fe5f89 100644 --- a/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py +++ b/tests/e2e/quota_management/spend_tracking/spend_e2e_client.py @@ -359,6 +359,17 @@ class SpendClient: def probe(self, path: str, *, params: DateRangeParams) -> ProbeResult: return self.proxy.transport.probe(path, params=params) + def probe_until_healthy(self, path: str, *, params: DateRangeParams) -> ProbeResult: + outcome: Final = await_converged( + lambda: self.probe(path, params=params), + converged=lambda result: result.healthy, + timeout=self.proxy.poll_timeout, + interval=self.proxy.poll_interval, + now=time.monotonic, + sleep=time.sleep, + ) + return outcome.result if isinstance(outcome, Converged) else outcome.last_result + def create_user(self, *, email: str, role: UserRole, user_id: str) -> str: return unwrap( self.proxy.transport.post( diff --git a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py index 8cb3e3927f0..67fd88bc84d 100644 --- a/tests/e2e/quota_management/spend_tracking/test_spend_routes.py +++ b/tests/e2e/quota_management/spend_tracking/test_spend_routes.py @@ -17,9 +17,11 @@ fast: no batch-write wait, no provider calls. """ from datetime import datetime, timedelta, timezone +from typing import Final import pytest +from e2e_http import ProbeResult from models import DateRangeParams from spend_e2e_client import SpendClient @@ -72,15 +74,10 @@ SPEND_ROUTES = ( _SPEND_PREFIXES = ("/spend", "/global/spend", "/global/activity") -_MISSING_VIEW_SKIP = pytest.mark.skip( - reason=( - "LIT-5211: on a fresh database the proxy's startup view creation can lose the race " - "against schema migrations, leaving MonthlyGlobalSpend/DailyTagSpend/Last30d* views " - "missing and these routes 500ing until the views exist" - ) -) - -_VIEW_BACKED_ROUTES = frozenset( +# Served from the MonthlyGlobalSpend / DailyTagSpend / Last30d* views, which the +# proxy creates in the background once the schema migrations have landed, so on a +# fresh database they can 500 for a while after the proxy starts serving. +_VIEW_BACKED_ROUTES: Final = frozenset( ( "/global/spend", "/global/spend/keys", @@ -98,15 +95,15 @@ def _date_range() -> DateRangeParams: return DateRangeParams(start_date=start.isoformat(), end_date=end.isoformat()) -@pytest.mark.parametrize( - "route", - tuple( - pytest.param(route, marks=_MISSING_VIEW_SKIP) if route in _VIEW_BACKED_ROUTES else route - for route in SPEND_ROUTES - ), -) +def _probe(client: SpendClient, route: str) -> ProbeResult: + if route in _VIEW_BACKED_ROUTES: + return client.probe_until_healthy(route, params=_date_range()) + return client.probe(route, params=_date_range()) + + +@pytest.mark.parametrize("route", SPEND_ROUTES) def test_spend_route_responsive(client: SpendClient, route: str) -> None: - result = client.probe(route, params=_date_range()) + result = _probe(client, route) print(f"{route} -> {result.status_code}\n{result.body[:600]}") assert result.healthy, f"{route} -> {result.status_code}\n{result.body[:600]}" 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 468180ffe52..f71f9c20f3b 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1810,6 +1810,37 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke assert announced == [] +class _ShutdownAwarePrisma(MockPrisma): + def __init__(self): + super().__init__() + self.stop_view_setup_task = AsyncMock() + + +@pytest.mark.asyncio +async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): + import yaml + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + fake_prisma = _ShutdownAwarePrisma() + config_path = tmp_path / "config.yaml" + with open(config_path, "w") as f: + yaml.dump({"general_settings": {"master_key": "sk-12345"}}, f) + monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path)) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake_prisma) + monkeypatch.setattr(proxy_server_module, "store_model_in_db", False) + + async with proxy_startup_event(FastAPI()): + stopped_while_serving = fake_prisma.stop_view_setup_task.await_count + + actual = { + "stopped_while_serving": stopped_while_serving, + "stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count, + } + assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1} + + def test_team_info_masking(): """ Test that sensitive team information is properly masked @@ -13608,6 +13639,7 @@ def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client.db.start_token_refresh_task = AsyncMock() client.check_view_exists = AsyncMock() client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + client.start_view_setup_task = MagicMock() client.start_db_health_watchdog_task = AsyncMock() client.health_check = AsyncMock(side_effect=health_check_error) return client @@ -13669,13 +13701,35 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) call_order = MagicMock() + call_order.attach_mock(mock_client.start_view_setup_task, "view_setup") call_order.attach_mock(mock_client.start_db_health_watchdog_task, "watchdog") call_order.attach_mock(mock_client.health_check, "health_check") await _run_setup_prisma_client(mock_client) assert mock_client.start_db_health_watchdog_task.await_count == 1 - assert [call[0] for call in call_order.mock_calls] == ["watchdog", "health_check"] + assert [call[0] for call in call_order.mock_calls] == ["view_setup", "watchdog", "health_check"] + + +@pytest.mark.asyncio +async def test_setup_prisma_client_hands_view_creation_to_the_held_task(monkeypatch): + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "True") + + mock_client = _mock_startup_prisma_client() + result = await _run_setup_prisma_client(mock_client) + + actual = { + "result": result, + "view_setup_started": mock_client.start_view_setup_task.call_count, + "direct_view_creation": mock_client.check_view_exists.await_count, + "direct_row_count": mock_client._set_spend_logs_row_count_in_proxy_state.await_count, + } + assert actual == { + "result": mock_client, + "view_setup_started": 1, + "direct_view_creation": 0, + "direct_row_count": 0, + } @pytest.mark.asyncio 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 18b02ac7772..9aa57c7a19b 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 @@ -5,11 +5,15 @@ Symbols pinned here: - ``PrismaClient.writer_db`` - ``PrismaClient.connect`` - ``PrismaClient.disconnect`` + - ``PrismaClient.start_view_setup_task`` + - ``PrismaClient.stop_view_setup_task`` + - ``PrismaClient._run_view_setup`` """ from __future__ import annotations import asyncio +import logging from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -17,6 +21,27 @@ import pytest from litellm.proxy.utils import PrismaClient +_PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present" + + +def _absent() -> list[dict[str, bool]]: + return [{"present": False}] + + +def _present() -> list[dict[str, bool]]: + return [{"present": True}] + + +def _wire_view_setup(prisma_client: PrismaClient, probe: AsyncMock) -> MagicMock: + prisma_client.db.query_raw = probe + prisma_client.check_view_exists = AsyncMock() + prisma_client._set_spend_logs_row_count_in_proxy_state = AsyncMock() + call_order = MagicMock() + call_order.attach_mock(probe, "probe") + call_order.attach_mock(prisma_client.check_view_exists, "views") + call_order.attach_mock(prisma_client._set_spend_logs_row_count_in_proxy_state, "row_count") + return call_order + @pytest.mark.asyncio async def test_prismaclient_init_wires_default_config( @@ -205,3 +230,221 @@ async def test_disconnect_raises_when_underlying_fails( prisma_client.db.disconnect = AsyncMock(side_effect=RuntimeError("disconnect boom")) with pytest.raises(RuntimeError, match="disconnect boom"): await prisma_client.disconnect() + + +@pytest.mark.asyncio +async def test_view_setup_waits_for_the_spend_logs_table_before_creating_views(prisma_client: PrismaClient) -> None: + probe = AsyncMock(side_effect=[_absent(), _absent(), _present()]) + call_order = _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + "probe_args": probe.await_args.args, + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "probe", "probe", "row_count", "views"], + "probe_args": (_PROBE_SQL, '"LiteLLM_SpendLogs"'), + } + + +@pytest.mark.asyncio +async def test_view_setup_probe_resolves_through_the_connection_search_path( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DATABASE_SCHEMA", "litellm_tenant") + probe = AsyncMock(return_value=_present()) + _wire_view_setup(prisma_client, probe) + + await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + assert probe.await_args.args == (_PROBE_SQL, '"LiteLLM_SpendLogs"') + + +@pytest.mark.asyncio +async def test_view_setup_sets_the_row_count_even_when_view_creation_keeps_failing( + prisma_client: PrismaClient, +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02) + + actual = { + "outcome": outcome, + "row_count_set": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count >= 1, + } + assert actual == {"outcome": "timed_out", "row_count_set": True} + + +@pytest.mark.asyncio +async def test_view_setup_gives_up_when_the_table_never_appears(prisma_client: PrismaClient) -> None: + probe = AsyncMock(return_value=_absent()) + _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.02) + + actual = { + "outcome": outcome, + "kept_polling": probe.await_count > 1, + "views_attempted": prisma_client.check_view_exists.await_count, + "row_count_attempted": prisma_client._set_spend_logs_row_count_in_proxy_state.await_count, + } + assert actual == { + "outcome": "timed_out", + "kept_polling": True, + "views_attempted": 0, + "row_count_attempted": 0, + } + + +@pytest.mark.asyncio +async def test_view_setup_retries_when_view_creation_fails_mid_migration(prisma_client: PrismaClient) -> None: + 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] + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "row_count", "views", "probe", "row_count", "views"], + } + + +@pytest.mark.asyncio +async def test_view_setup_retries_when_the_table_probe_itself_fails(prisma_client: PrismaClient) -> None: + probe = AsyncMock(side_effect=[RuntimeError("connection reset"), _present()]) + call_order = _wire_view_setup(prisma_client, probe) + + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=5) + + actual = { + "outcome": outcome, + "calls": [call[0] for call in call_order.mock_calls], + } + assert actual == { + "outcome": "ready", + "calls": ["probe", "probe", "row_count", "views"], + } + + +@pytest.mark.asyncio +async def test_run_view_setup_logs_an_error_naming_the_table_on_timeout( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + actual = { + "outcome": outcome, + "error_count": len(errors), + "names_table": '"LiteLLM_SpendLogs"' in errors[0], + "tells_operator_to_migrate": "migrations" in errors[0] and "restart" in errors[0], + } + assert actual == { + "outcome": "timed_out", + "error_count": 1, + "names_table": True, + "tells_operator_to_migrate": True, + } + + +@pytest.mark.asyncio +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: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + prisma_client.check_view_exists.side_effect = RuntimeError("permission denied for schema public") + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + errors = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + actual = { + "outcome": outcome, + "error_count": len(errors), + "names_the_error": "permission denied for schema public" in errors[0], + "blames_missing_migrations": "did not appear" in errors[0], + "tells_operator_to_restart": "restart" in errors[0], + } + assert actual == { + "outcome": "timed_out", + "error_count": 1, + "names_the_error": True, + "blames_missing_migrations": False, + "tells_operator_to_restart": True, + } + + +@pytest.mark.asyncio +async def test_run_view_setup_stays_quiet_when_views_are_ready( + prisma_client: PrismaClient, caplog: pytest.LogCaptureFixture +) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_present())) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + outcome = await prisma_client._run_view_setup(poll_interval_seconds=0.001, deadline_seconds=0.01) + + actual = { + "outcome": outcome, + "errors": [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR], + } + assert actual == {"outcome": "ready", "errors": []} + + +@pytest.mark.asyncio +async def test_stop_view_setup_task_cancels_a_task_parked_between_polls(prisma_client: PrismaClient) -> None: + probe = AsyncMock(return_value=_absent()) + _wire_view_setup(prisma_client, probe) + + prisma_client.start_view_setup_task() + task = prisma_client._view_setup_task + await asyncio.sleep(0) + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + + actual = { + "probed_before_parking": probe.await_count, + "task_cancelled": task is not None and task.cancelled(), + "reference_cleared": prisma_client._view_setup_task, + "views_attempted": prisma_client.check_view_exists.await_count, + } + assert actual == { + "probed_before_parking": 1, + "task_cancelled": True, + "reference_cleared": None, + "views_attempted": 0, + } + + +@pytest.mark.asyncio +async def test_stop_view_setup_task_is_a_noop_without_a_task(prisma_client: PrismaClient) -> None: + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + assert prisma_client._view_setup_task is None + + +@pytest.mark.asyncio +async def test_start_view_setup_task_twice_keeps_the_first_task(prisma_client: PrismaClient) -> None: + _wire_view_setup(prisma_client, AsyncMock(return_value=_absent())) + + prisma_client.start_view_setup_task() + first = prisma_client._view_setup_task + prisma_client.start_view_setup_task() + second = prisma_client._view_setup_task + await asyncio.wait_for(prisma_client.stop_view_setup_task(), timeout=1) + + actual = { + "first_is_task": isinstance(first, asyncio.Task), + "second_is_first": second is first, + } + assert actual == {"first_is_task": True, "second_is_first": True}