diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e15b48577de..30ace30fc47 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1965,9 +1965,9 @@ class ProxyLogging: normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type - litellm_logging_obj.model_call_details[ - "call_type" - ] = normalized_call_type + litellm_logging_obj.model_call_details["call_type"] = ( + normalized_call_type + ) # Pass-through endpoints are logged via the callback loop's # async_post_call_failure_hook — skip pre_call and failure handlers. if litellm_logging_obj.call_type == CallTypes.pass_through.value: @@ -2065,10 +2065,12 @@ class ProxyLogging: raise else: try: - guardrail_response = await callback.async_post_call_success_hook( - user_api_key_dict=user_api_key_dict, - data=data, - response=response, + guardrail_response = ( + await callback.async_post_call_success_hook( + user_api_key_dict=user_api_key_dict, + data=data, + response=response, + ) ) except Exception as e: _enrich_http_exception_with_guardrail_context(e, callback) @@ -2281,13 +2283,15 @@ class ProxyLogging: "async_post_call_streaming_iterator_hook" in type(callback).__dict__ ): - current_response = self._wrap_streaming_iterator_with_enrichment( - _callback, - _callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=current_response, - request_data=request_data, - ), + current_response = ( + self._wrap_streaming_iterator_with_enrichment( + _callback, + _callback.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=current_response, + request_data=request_data, + ), + ) ) elif "apply_guardrail" in type(callback).__dict__: request_data["guardrail_to_apply"] = callback @@ -2300,13 +2304,15 @@ class ProxyLogging: ), ) else: - current_response = self._wrap_streaming_iterator_with_enrichment( - _callback, - _callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=current_response, - request_data=request_data, - ), + current_response = ( + self._wrap_streaming_iterator_with_enrichment( + _callback, + _callback.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=current_response, + request_data=request_data, + ), + ) ) # Actually iterate through the chained async generator and yield chunks @@ -2512,6 +2518,7 @@ class PrismaClient: self._watching_engine: bool = False self._engine_confirmed_dead: bool = False self._engine_wait_thread: Optional[threading.Thread] = None + self._health_check_last_cleanup_ts: float = 0.0 verbose_proxy_logger.debug("Success - Created Prisma Client") def get_request_status( @@ -4487,7 +4494,17 @@ class PrismaClient: ) verbose_proxy_logger.debug(f"Saving health check data: {health_check_data}") - return await self.db.litellm_healthchecktable.create(data=health_check_data) + result = await self.db.litellm_healthchecktable.create( + data=health_check_data + ) + + # Prune old rows at most once per hour to prevent unbounded table growth + now = time.time() + if now - self._health_check_last_cleanup_ts >= 3600: + self._health_check_last_cleanup_ts = now + await self.cleanup_old_health_checks() + + return result except Exception as e: verbose_proxy_logger.error( @@ -4525,25 +4542,29 @@ class PrismaClient: async def get_all_latest_health_checks(self): """ - Get the latest health check for each model + Get the latest health check for each model. + + Only considers rows within the retention window (HEALTH_CHECK_TTL_DAYS) to + avoid full-table scans on deployments with large health check histories. """ try: - # Get all unique model names first + health_check_ttl_days = int(os.getenv("HEALTH_CHECK_TTL_DAYS", "7")) + cutoff = datetime.utcnow() - timedelta(days=health_check_ttl_days) all_checks = await self.db.litellm_healthchecktable.find_many( - order={"checked_at": "desc"} + where={"checked_at": {"gte": cutoff}}, + order={"checked_at": "desc"}, ) - # Group by model_name and get the latest for each + # Group by model and return only the latest entry per model latest_checks = {} for check in all_checks: - # Create a unique key: prefer model_id if available, otherwise use model_name - # This ensures we get the latest check for each unique model + # Prefer model_id as the dedup key; fall back to model_name if check.model_id: key = (check.model_id, check.model_name) else: key = (None, check.model_name) - # Only add if we haven't seen this key yet (since checks are ordered by checked_at desc) + # First entry wins because rows are ordered by checked_at desc if key not in latest_checks: latest_checks[key] = check @@ -4552,6 +4573,25 @@ class PrismaClient: verbose_proxy_logger.error(f"Error getting all latest health checks: {e}") return [] + async def cleanup_old_health_checks(self) -> None: + """ + Delete LiteLLM_HealthCheckTable rows older than HEALTH_CHECK_TTL_DAYS (default 7). + + Called from save_health_check_result at most once per hour to prevent + unbounded table growth without adding per-insert overhead. + """ + try: + health_check_ttl_days = int(os.getenv("HEALTH_CHECK_TTL_DAYS", "7")) + cutoff = datetime.utcnow() - timedelta(days=health_check_ttl_days) + result = await self.db.litellm_healthchecktable.delete_many( + where={"checked_at": {"lt": cutoff}} + ) + verbose_proxy_logger.debug( + f"Health check cleanup: deleted {result} rows older than {health_check_ttl_days} days" + ) + except Exception as e: + verbose_proxy_logger.error(f"Error cleaning up old health checks: {e}") + ### HELPER FUNCTIONS ### @@ -4947,25 +4987,27 @@ async def update_daily_tag_spend( ): """ Separate scheduler job to commit daily tag spend updates. - + Runs at a longer interval (2.3x default) than the main update_spend job to reduce query contention for DailyTagSpend table. - + This is called by a dedicated scheduler job and does NOT process: - Regular spend updates (user, key, team, org) - End-user spend - Agent spend - Spend logs - + Only processes tag spend transactions from the daily_tag_spend_update_queue. - + Args: prisma_client: PrismaClient instance proxy_logging_obj: ProxyLogging instance for error handling """ n_retry_times = 3 try: - if proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis(): + if ( + proxy_logging_obj.db_spend_update_writer.redis_update_buffer._should_commit_spend_updates_to_redis() + ): await proxy_logging_obj.db_spend_update_writer._commit_daily_tag_spend_to_db_with_redis( prisma_client=prisma_client, n_retry_times=n_retry_times, diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index 13d2131efad..806dd619de3 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -24,33 +24,53 @@ from litellm.proxy.utils import PrismaClient def mock_prisma(): """Simplified mock PrismaClient with bound methods""" client = MagicMock() - client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "test-id"}) - client.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[{"id": "1", "model_name": "test"}]) - + client.db.litellm_healthchecktable.create = AsyncMock( + return_value={"id": "test-id"} + ) + client.db.litellm_healthchecktable.find_many = AsyncMock( + return_value=[{"id": "1", "model_name": "test"}] + ) + # Bind actual methods import types - for method in ['save_health_check_result', '_validate_response_time', '_clean_details', - 'get_health_check_history', 'get_all_latest_health_checks']: + + for method in [ + "save_health_check_result", + "_validate_response_time", + "_clean_details", + "get_health_check_history", + "get_all_latest_health_checks", + ]: setattr(client, method, types.MethodType(getattr(PrismaClient, method), client)) - + return client @pytest.mark.asyncio -@pytest.mark.parametrize("status,healthy,unhealthy,should_succeed", [ - ("healthy", 1, 0, True), - ("unhealthy", 0, 1, True), - ("healthy", 1, 0, False), # Database error case -]) -async def test_save_health_check_result(mock_prisma, status, healthy, unhealthy, should_succeed): +@pytest.mark.parametrize( + "status,healthy,unhealthy,should_succeed", + [ + ("healthy", 1, 0, True), + ("unhealthy", 0, 1, True), + ("healthy", 1, 0, False), # Database error case + ], +) +async def test_save_health_check_result( + mock_prisma, status, healthy, unhealthy, should_succeed +): """Test health check result saving with various scenarios""" if not should_succeed: - mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception("DB Error") - + mock_prisma.db.litellm_healthchecktable.create.side_effect = Exception( + "DB Error" + ) + result = await mock_prisma.save_health_check_result( - model_name="test-model", status=status, healthy_count=healthy, unhealthy_count=unhealthy + model_name="test-model", + status=status, + healthy_count=healthy, + unhealthy_count=unhealthy, ) - + if should_succeed: mock_prisma.db.litellm_healthchecktable.create.assert_called_once() else: @@ -66,24 +86,31 @@ async def test_get_health_check_history(mock_prisma): @pytest.mark.asyncio -@pytest.mark.parametrize("healthy_count,unhealthy_count,expected_status", [ - (1, 0, "healthy"), - (0, 1, "unhealthy"), - (2, 1, "healthy"), -]) +@pytest.mark.parametrize( + "healthy_count,unhealthy_count,expected_status", + [ + (1, 0, "healthy"), + (0, 1, "unhealthy"), + (2, 1, "healthy"), + ], +) async def test_save_health_check_to_db(healthy_count, unhealthy_count, expected_status): """Test _save_health_check_to_db function with different endpoint counts""" mock_client = MagicMock() mock_client.save_health_check_result = AsyncMock() - + healthy_endpoints = [{"model": "test"}] * healthy_count unhealthy_endpoints = [{"error": "test error"}] * unhealthy_count - + await _save_health_check_to_db( - mock_client, "test-model", healthy_endpoints, unhealthy_endpoints, - 1234567890.0, "test-user" + mock_client, + "test-model", + healthy_endpoints, + unhealthy_endpoints, + 1234567890.0, + "test-user", ) - + call_args = mock_client.save_health_check_result.call_args[1] assert call_args["status"] == expected_status assert call_args["healthy_count"] == healthy_count @@ -99,6 +126,7 @@ async def test_save_health_check_to_db_no_client(): # Tests for background health check functions + def test_build_model_param_to_info_mapping(): """Test building model parameter to info mapping""" model_list = [ @@ -118,9 +146,9 @@ def test_build_model_param_to_info_mapping(): "litellm_params": {"model": "gpt-3.5-turbo"}, # Same model param }, ] - + result = _build_model_param_to_info_mapping(model_list) - + assert "gpt-3.5-turbo" in result assert "gpt-4" in result assert len(result["gpt-3.5-turbo"]) == 2 # Two models share same param @@ -139,7 +167,7 @@ def test_build_model_param_to_info_mapping_no_model_name(): "litellm_params": {"model": "gpt-3.5-turbo"}, }, ] - + result = _build_model_param_to_info_mapping(model_list) assert len(result) == 0 @@ -154,25 +182,25 @@ def test_aggregate_health_check_results(): {"model_name": "gpt-4", "model_id": "model-456"}, ], } - + healthy_endpoints = [ {"model": "gpt-3.5-turbo"}, ] unhealthy_endpoints = [ {"model": "gpt-4", "error": "Rate limit exceeded"}, ] - + result = _aggregate_health_check_results( model_param_to_info, healthy_endpoints, unhealthy_endpoints ) - + # Check gpt-3.5-turbo is healthy gpt35_key = ("model-123", "gpt-3.5-turbo") assert gpt35_key in result assert result[gpt35_key]["healthy_count"] == 1 assert result[gpt35_key]["unhealthy_count"] == 0 assert result[gpt35_key]["error_message"] is None - + # Check gpt-4 is unhealthy gpt4_key = ("model-456", "gpt-4") assert gpt4_key in result @@ -188,17 +216,17 @@ def test_aggregate_health_check_results_multiple_endpoints(): {"model_name": "gpt-3.5-turbo", "model_id": "model-123"}, ], } - + healthy_endpoints = [ {"model": "gpt-3.5-turbo"}, {"model": "gpt-3.5-turbo"}, ] unhealthy_endpoints = [] - + result = _aggregate_health_check_results( model_param_to_info, healthy_endpoints, unhealthy_endpoints ) - + key = ("model-123", "gpt-3.5-turbo") assert result[key]["healthy_count"] == 2 assert result[key]["unhealthy_count"] == 0 @@ -209,7 +237,7 @@ async def test_save_health_check_results_if_changed_status_changed(): """Test saving when status changes""" mock_prisma = MagicMock() mock_prisma.save_health_check_result = AsyncMock() - + model_results = { ("model-123", "gpt-3.5-turbo"): { "model_name": "gpt-3.5-turbo", @@ -219,7 +247,7 @@ async def test_save_health_check_results_if_changed_status_changed(): "error_message": None, }, } - + # Latest check shows unhealthy, new result is healthy (status changed) latest_checks_map = { "model-123": MagicMock( @@ -227,12 +255,16 @@ async def test_save_health_check_results_if_changed_status_changed(): checked_at=datetime.now(timezone.utc) - timedelta(minutes=5), ), } - + start_time = 1234567890.0 await _save_health_check_results_if_changed( - mock_prisma, model_results, latest_checks_map, start_time, "background_health_check" + mock_prisma, + model_results, + latest_checks_map, + start_time, + "background_health_check", ) - + # Should save because status changed mock_prisma.save_health_check_result.assert_called_once() call_kwargs = mock_prisma.save_health_check_result.call_args[1] @@ -246,7 +278,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_recent(): """Test skipping save when status unchanged and checked recently""" mock_prisma = MagicMock() mock_prisma.save_health_check_result = AsyncMock() - + model_results = { ("model-123", "gpt-3.5-turbo"): { "model_name": "gpt-3.5-turbo", @@ -256,7 +288,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_recent(): "error_message": None, }, } - + # Latest check shows healthy, new result is healthy (status unchanged) # And checked recently (within 1 hour) latest_checks_map = { @@ -265,12 +297,16 @@ async def test_save_health_check_results_if_changed_status_unchanged_recent(): checked_at=datetime.now(timezone.utc) - timedelta(minutes=30), ), } - + start_time = 1234567890.0 await _save_health_check_results_if_changed( - mock_prisma, model_results, latest_checks_map, start_time, "background_health_check" + mock_prisma, + model_results, + latest_checks_map, + start_time, + "background_health_check", ) - + # Should NOT save because status unchanged and checked recently mock_prisma.save_health_check_result.assert_not_called() @@ -280,7 +316,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_old(): """Test saving when status unchanged but last check is old (>1 hour)""" mock_prisma = MagicMock() mock_prisma.save_health_check_result = AsyncMock() - + model_results = { ("model-123", "gpt-3.5-turbo"): { "model_name": "gpt-3.5-turbo", @@ -290,7 +326,7 @@ async def test_save_health_check_results_if_changed_status_unchanged_old(): "error_message": None, }, } - + # Latest check shows healthy, new result is healthy (status unchanged) # But checked >1 hour ago latest_checks_map = { @@ -299,12 +335,16 @@ async def test_save_health_check_results_if_changed_status_unchanged_old(): checked_at=datetime.now(timezone.utc) - timedelta(hours=2), ), } - + start_time = 1234567890.0 await _save_health_check_results_if_changed( - mock_prisma, model_results, latest_checks_map, start_time, "background_health_check" + mock_prisma, + model_results, + latest_checks_map, + start_time, + "background_health_check", ) - + # Should save because last check is old (>1 hour) mock_prisma.save_health_check_result.assert_called_once() @@ -314,7 +354,7 @@ async def test_save_health_check_results_if_changed_no_previous_check(): """Test saving when there's no previous check""" mock_prisma = MagicMock() mock_prisma.save_health_check_result = AsyncMock() - + model_results = { ("model-123", "gpt-3.5-turbo"): { "model_name": "gpt-3.5-turbo", @@ -324,15 +364,19 @@ async def test_save_health_check_results_if_changed_no_previous_check(): "error_message": None, }, } - + # No previous check latest_checks_map = {} - + start_time = 1234567890.0 await _save_health_check_results_if_changed( - mock_prisma, model_results, latest_checks_map, start_time, "background_health_check" + mock_prisma, + model_results, + latest_checks_map, + start_time, + "background_health_check", ) - + # Should save because no previous check mock_prisma.save_health_check_result.assert_called_once() @@ -343,7 +387,7 @@ async def test_save_background_health_checks_to_db(): mock_prisma = MagicMock() mock_prisma.save_health_check_result = AsyncMock() mock_prisma.get_all_latest_health_checks = AsyncMock(return_value=[]) - + model_list = [ { "model_name": "gpt-3.5-turbo", @@ -351,20 +395,25 @@ async def test_save_background_health_checks_to_db(): "litellm_params": {"model": "gpt-3.5-turbo"}, }, ] - + healthy_endpoints = [{"model": "gpt-3.5-turbo"}] unhealthy_endpoints = [] - + start_time = 1234567890.0 - + await _save_background_health_checks_to_db( - mock_prisma, model_list, healthy_endpoints, unhealthy_endpoints, start_time, "background_health_check" + mock_prisma, + model_list, + healthy_endpoints, + unhealthy_endpoints, + start_time, + "background_health_check", ) - + # Should call get_all_latest_health_checks and save_health_check_result mock_prisma.get_all_latest_health_checks.assert_called_once() mock_prisma.save_health_check_result.assert_called_once() - + call_kwargs = mock_prisma.save_health_check_result.call_args[1] assert call_kwargs["model_name"] == "gpt-3.5-turbo" assert call_kwargs["model_id"] == "model-123" @@ -385,8 +434,10 @@ async def test_save_background_health_checks_to_db_no_prisma(): async def test_save_background_health_checks_to_db_exception_handling(): """Test exception handling in background health check save""" mock_prisma = MagicMock() - mock_prisma.get_all_latest_health_checks = AsyncMock(side_effect=Exception("DB Error")) - + mock_prisma.get_all_latest_health_checks = AsyncMock( + side_effect=Exception("DB Error") + ) + model_list = [ { "model_name": "gpt-3.5-turbo", @@ -394,12 +445,12 @@ async def test_save_background_health_checks_to_db_exception_handling(): "litellm_params": {"model": "gpt-3.5-turbo"}, }, ] - + # Should not raise exception, should handle gracefully await _save_background_health_checks_to_db( mock_prisma, model_list, [], [], 0.0, "background_health_check" ) - + # Function should complete without raising @@ -411,32 +462,34 @@ async def test_get_all_latest_health_checks_with_model_id(mock_prisma): mock_check1.model_id = "model-123" mock_check1.model_name = "gpt-3.5-turbo" mock_check1.checked_at = datetime.now(timezone.utc) - timedelta(minutes=10) - + mock_check2 = MagicMock() mock_check2.model_id = "model-456" mock_check2.model_name = "gpt-3.5-turbo" mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=5) - + mock_check3 = MagicMock() mock_check3.model_id = "model-123" mock_check3.model_name = "gpt-3.5-turbo" - mock_check3.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest for model-123 - + mock_check3.checked_at = datetime.now(timezone.utc) - timedelta( + minutes=1 + ) # Latest for model-123 + # Order by checked_at desc mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( return_value=[mock_check3, mock_check2, mock_check1] ) - + result = await mock_prisma.get_all_latest_health_checks() - + # Should return 2 unique models (by model_id) assert len(result) == 2 - + # Should have latest check for each model_id model_ids = {check.model_id for check in result} assert "model-123" in model_ids assert "model-456" in model_ids - + # model-123 should have the latest check (1 minute ago) model123_check = next(c for c in result if c.model_id == "model-123") assert model123_check.checked_at == mock_check3.checked_at @@ -449,18 +502,18 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma): mock_check1.model_id = None mock_check1.model_name = "gpt-3.5-turbo" mock_check1.checked_at = datetime.now(timezone.utc) - timedelta(minutes=10) - + mock_check2 = MagicMock() mock_check2.model_id = None mock_check2.model_name = "gpt-3.5-turbo" mock_check2.checked_at = datetime.now(timezone.utc) - timedelta(minutes=1) # Latest - + mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock( return_value=[mock_check2, mock_check1] ) - + result = await mock_prisma.get_all_latest_health_checks() - + # Should return 1 unique model (by model_name) assert len(result) == 1 assert result[0].model_name == "gpt-3.5-turbo" @@ -480,7 +533,14 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c healthy = [{"model": "gpt-4"}] unhealthy = [] - async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None, max_concurrency=None): + async def mock_perform_health_check( + model_list, + model=None, + cli_model=None, + details=True, + model_id=None, + max_concurrency=None, + ): return healthy, unhealthy, {} with patch( @@ -505,5 +565,112 @@ async def test_perform_health_check_and_save_passes_model_id_to_perform_health_c assert result["unhealthy_count"] == 0 +@pytest.mark.asyncio +async def test_get_all_latest_health_checks_applies_ttl_filter(mock_prisma): + """get_all_latest_health_checks should pass a checked_at >= cutoff filter so + it does not load the entire table into memory.""" + mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[]) + + await mock_prisma.get_all_latest_health_checks() + + call_kwargs = mock_prisma.db.litellm_healthchecktable.find_many.call_args[1] + assert "where" in call_kwargs + assert "checked_at" in call_kwargs["where"] + assert "gte" in call_kwargs["where"]["checked_at"] + + +@pytest.mark.asyncio +async def test_get_all_latest_health_checks_ttl_env_override(mock_prisma): + """HEALTH_CHECK_TTL_DAYS env var should control the cutoff window.""" + mock_prisma.db.litellm_healthchecktable.find_many = AsyncMock(return_value=[]) + + with patch.dict(os.environ, {"HEALTH_CHECK_TTL_DAYS": "3"}): + await mock_prisma.get_all_latest_health_checks() + + call_kwargs = mock_prisma.db.litellm_healthchecktable.find_many.call_args[1] + cutoff = call_kwargs["where"]["checked_at"]["gte"] + # The cutoff should be roughly 3 days ago (within a small tolerance) + expected = datetime.utcnow() - timedelta(days=3) + assert abs((expected - cutoff).total_seconds()) < 5 + + +@pytest.mark.asyncio +async def test_cleanup_old_health_checks_calls_delete_many(mock_prisma): + """cleanup_old_health_checks should delete rows older than TTL via delete_many.""" + import types + + mock_prisma.db.litellm_healthchecktable.delete_many = AsyncMock(return_value=42) + mock_prisma.cleanup_old_health_checks = types.MethodType( + PrismaClient.cleanup_old_health_checks, mock_prisma + ) + + await mock_prisma.cleanup_old_health_checks() + + call_kwargs = mock_prisma.db.litellm_healthchecktable.delete_many.call_args[1] + assert "where" in call_kwargs + assert "checked_at" in call_kwargs["where"] + assert "lt" in call_kwargs["where"]["checked_at"] + + +@pytest.mark.asyncio +async def test_save_health_check_result_triggers_cleanup_after_one_hour(): + """cleanup_old_health_checks should be called when last cleanup was > 1 hour ago.""" + import types + + mock_client = MagicMock() + mock_client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "x"}) + mock_client.db.litellm_healthchecktable.delete_many = AsyncMock(return_value=0) + mock_client._health_check_last_cleanup_ts = 0.0 # never cleaned up + # Bind helper methods properly so `self` is resolved correctly + mock_client._validate_response_time = types.MethodType( + PrismaClient._validate_response_time, mock_client + ) + mock_client._clean_details = types.MethodType( + PrismaClient._clean_details, mock_client + ) + mock_client.save_health_check_result = types.MethodType( + PrismaClient.save_health_check_result, mock_client + ) + mock_client.cleanup_old_health_checks = types.MethodType( + PrismaClient.cleanup_old_health_checks, mock_client + ) + + await mock_client.save_health_check_result( + model_name="gpt-4", status="healthy", healthy_count=1, unhealthy_count=0 + ) + + mock_client.db.litellm_healthchecktable.delete_many.assert_called_once() + + +@pytest.mark.asyncio +async def test_save_health_check_result_skips_cleanup_within_one_hour(): + """cleanup_old_health_checks should NOT be called if last cleanup was < 1 hour ago.""" + import types + + mock_client = MagicMock() + mock_client.db.litellm_healthchecktable.create = AsyncMock(return_value={"id": "x"}) + mock_client.db.litellm_healthchecktable.delete_many = AsyncMock(return_value=0) + mock_client._health_check_last_cleanup_ts = time.time() # just ran + # Bind helper methods properly so `self` is resolved correctly + mock_client._validate_response_time = types.MethodType( + PrismaClient._validate_response_time, mock_client + ) + mock_client._clean_details = types.MethodType( + PrismaClient._clean_details, mock_client + ) + mock_client.save_health_check_result = types.MethodType( + PrismaClient.save_health_check_result, mock_client + ) + mock_client.cleanup_old_health_checks = types.MethodType( + PrismaClient.cleanup_old_health_checks, mock_client + ) + + await mock_client.save_health_check_result( + model_name="gpt-4", status="healthy", healthy_count=1, unhealthy_count=0 + ) + + mock_client.db.litellm_healthchecktable.delete_many.assert_not_called() + + if __name__ == "__main__": - pytest.main([__file__]) \ No newline at end of file + pytest.main([__file__])