diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index eaab34adec..14b49901c9 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2205,6 +2205,7 @@ async def test_proxy_server_prisma_setup(): 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() diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 1d9bb15b21..cc1cc27839 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1623,6 +1623,7 @@ async def test_health_check_not_called_when_disabled(monkeypatch): mock_prisma.health_check = AsyncMock() mock_prisma.check_view_exists = AsyncMock() mock_prisma._set_spend_logs_row_count_in_proxy_state = AsyncMock() + mock_prisma.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()