test(gcs): cover the requeue drop on a refilled queue and the healthy GCS flush in /key/health

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

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