diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py index 85306a10428..cee457fab50 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket.py @@ -19,6 +19,7 @@ class _FakeUploadGCSLogger(GCSBucketLogger): super().__init__(bucket_name="test-bucket") self.log_queue = asyncio.Queue(maxsize=queue_maxsize) self.failing_ids: frozenset[str] = frozenset() + self.arriving_during_upload: tuple[str, ...] = () self.uploaded: list[list[str]] = [] async def enqueue(self, request_id: str) -> None: @@ -44,6 +45,8 @@ class _FakeUploadGCSLogger(GCSBucketLogger): if isinstance(logging_payload, str) else [logging_payload["id"]] ) + for request_id in self.arriving_during_upload: + await self.enqueue(request_id) if self.failing_ids.intersection(ids): raise RuntimeError("storage.googleapis.com returned 404") self.uploaded.append(ids) @@ -97,6 +100,20 @@ async def test_enqueue_on_a_full_queue_whose_flush_failed_drops_the_oldest_event assert logger.queued_ids() == ["req-2", "req-3"] +@pytest.mark.asyncio +async def test_failed_batch_is_dropped_when_new_events_filled_the_queue_during_the_upload(): + logger = _FakeUploadGCSLogger(queue_maxsize=2) + logger.failing_ids = frozenset({"req-1"}) + logger.arriving_during_upload = ("req-3", "req-4") + await logger.enqueue("req-1") + await logger.enqueue("req-2") + + result = await logger.flush_queue_and_report() + + assert result == GCSFlushResult(sent=0, failed=2) + assert logger.queued_ids() == ["req-3", "req-4"] + + @pytest.mark.asyncio async def test_empty_queue_flush_reports_nothing_sent_or_failed(): logger = _FakeUploadGCSLogger() diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e7e0e0e53fd..7251dd77d5b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -18441,8 +18441,7 @@ async def test_flush_gcs_reports_the_failed_upload_count_from_the_registered_log assert await flush_gcs_and_describe_failures(_gcs_logger_whose_flush_reports(sent=2, failed=0)) is None -@pytest.mark.asyncio -async def test_key_logging_marks_the_key_unhealthy_when_the_gcs_flush_leaves_events_undelivered(): +async def _key_logging_status_after_gcs_flush(sent: int, failed: int) -> LoggingCallbackStatus: from starlette.requests import Request as StarletteRequest from litellm.proxy.management_endpoints.key_management_endpoints import test_key_logging @@ -18455,15 +18454,29 @@ async def test_key_logging_marks_the_key_unhealthy_when_the_gcs_flush_leaves_eve patch("litellm.proxy.proxy_server.proxy_config", _default_team_gcs_proxy_config("team-gcs")), # test-quality-ok: test_key_logging reads the module-level proxy config patch( # test-quality-ok: the registered logger is a process-wide registry, not an injectable "litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class", - return_value=_gcs_logger_whose_flush_reports(sent=1, failed=3), + return_value=_gcs_logger_whose_flush_reports(sent=sent, failed=failed), ), ): - status = await test_key_logging(user_api_key_dict=caller, request=request, logging_callbacks=("gcs_bucket",)) + return await test_key_logging(user_api_key_dict=caller, request=request, logging_callbacks=("gcs_bucket",)) + + +@pytest.mark.asyncio +async def test_key_logging_marks_the_key_unhealthy_when_the_gcs_flush_leaves_events_undelivered(): + status = await _key_logging_status_after_gcs_flush(sent=1, failed=3) assert status["status"] == "unhealthy" assert "GCS upload failed for 3 event(s), 1 uploaded" in (status["details"] or "") +@pytest.mark.asyncio +async def test_key_logging_stays_healthy_when_the_gcs_flush_delivers_every_event(): + status = await _key_logging_status_after_gcs_flush(sent=1, failed=0) + + assert status["status"] == "healthy" + assert status["callbacks"] == ("gcs_bucket",) + assert "Manually check if logs were sent to gcs_bucket" in (status["details"] or "") + + @pytest.mark.asyncio async def test_flush_gcs_names_a_missing_logger_when_the_callback_never_initialized(): from litellm.proxy.management_endpoints.key_management_endpoints import flush_gcs_and_describe_failures