test(passthrough): cover streaming iterators routing through the proxy success handler

Part of LIT-9423

tests/unit/google_genai/test_google_genai_streaming_iterator.py

@@ -1,5 +1,6 @@
+import asyncio
 import json
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -8,6 +9,8 @@ from litellm.google_genai.streaming_iterator import (
     GoogleGenAIGenerateContentStreamingIterator,
 )
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.proxy.pass_through_endpoints import success_handler
+from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
 from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
 
 
@@ -154,3 +157,27 @@ async def test_async_streaming_iterator_forwards_sse_comment_events():
 
     chunk = await iterator.__anext__()
     assert chunk == b": keepalive\n\n"
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_logging_bills_through_the_proxy_pass_through_success_handler():  # test-quality-ok: no injection seam, the handler object identity passed to the router is the contract
+    iterator = AsyncGoogleGenAIGenerateContentStreamingIterator(
+        response=MagicMock(),
+        model="gemini-3.1-flash-image",
+        logging_obj=MagicMock(spec=LiteLLMLoggingObj),
+        generate_content_provider_config=MagicMock(),
+        litellm_metadata={},
+        custom_llm_provider="gemini",
+    )
+    iterator.collected_chunks.append(b"data: {}\n\n")
+
+    with patch.object(PassThroughStreamingHandler, "route_streaming_logging_to_handler", new=AsyncMock()) as route:
+        await iterator._handle_async_streaming_logging()
+        await asyncio.sleep(0)
+
+    route.assert_awaited_once()
+    assert route.await_args.kwargs["passthrough_success_handler_obj"] is (
+        success_handler.GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ
+    )
+    assert route.await_args.kwargs["endpoint_type"] is EndpointType.GEMINI
+    assert route.await_args.kwargs["raw_bytes"] == [b"data: {}\n\n"]

tests/unit/llms/anthropic/pass_through/messages/test_streaming_iterator.py

@@ -1,8 +1,9 @@
 import asyncio
 import json
 from datetime import datetime
-from unittest.mock import patch
+from unittest.mock import AsyncMock, patch
 
+import httpx
 import pytest
 
 
@@ -22,6 +23,8 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import (
     is_provider_error_chunk,
     parse_anthropic_error_event,
 )
+from litellm.proxy.pass_through_endpoints import success_handler
+from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler
 
 
 class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator):
@@ -1523,3 +1526,30 @@ async def test_relay_teardown_dispatches_deferred_billing_when_sentinel_never_co
     assert len(worker.enqueued) == 1, "teardown billing enqueued alongside the deferred dispatch"
     await worker.enqueued[0]
     assert deferred_fired.is_set()
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_response_iterator_relays_bytes_and_bills_through_proxy_success_handler():
+    body = b'event: message_stop\ndata: {"type": "message_stop"}\n\n'
+    response = httpx.Response(200, content=body)
+    logging_obj = _make_logging_obj("test_async_streaming_response_iterator_relays_bytes")
+    iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={})
+
+    with patch.object(PassThroughStreamingHandler, "route_streaming_logging_to_handler", new=AsyncMock()) as route:
+        relayed = [
+            chunk
+            async for chunk in iterator.get_async_streaming_response_iterator(
+                httpx_response=response, request_body={}, litellm_logging_obj=logging_obj
+            )
+        ]
+        for _ in range(100):
+            if route.await_count:
+                break
+            await asyncio.sleep(0.01)
+
+    assert b"".join(relayed) == body
+    route.assert_awaited_once()
+    assert route.await_args.kwargs["passthrough_success_handler_obj"] is (
+        success_handler.GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ
+    )
+    assert route.await_args.kwargs["url_route"] == "/v1/messages"