From 7818cd4aaf2d66ad0974bc6892c40cd9f2f2379a Mon Sep 17 00:00:00 2001
From: Bytechoreographer <Bytechoreographer@users.noreply.github.com>
Date: Fri, 26 Jun 2026 13:05:48 +0800
Subject: [PATCH] test(proxy): cover VERTEX_AI cost-injection branch in
 pass-through TTFT regression test

---
 .../test_streaming_handler_ttft.py            | 41 +++++++++++++++++++
 1 file changed, 41 insertions(+)

diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py
index 71b93904845..0af09c1d66c 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py
@@ -13,6 +13,7 @@ import httpx
 import pytest
 
 import litellm
+from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
 from litellm.proxy.pass_through_endpoints.streaming_handler import (
     PassThroughStreamingHandler,
 )
@@ -106,3 +107,43 @@ async def test_completion_start_time_stamped_on_cost_injection_path():
 
     assert received  # cost-injection branch was exercised
     assert isinstance(logging_obj.completion_start_time, datetime)
+
+
+@pytest.mark.asyncio
+async def test_completion_start_time_stamped_on_vertex_cost_injection_path():
+    response = _make_streaming_response([b'data: {"x": 1}\n\n'])
+    logging_obj = _logging_obj()
+
+    with (
+        patch.object(
+            PassThroughStreamingHandler,
+            "_route_streaming_logging_to_handler",
+            new=AsyncMock(),
+        ),
+        patch.object(litellm, "include_cost_in_streaming_usage", True),
+        patch.object(
+            PassThroughStreamingHandler,
+            "_extract_model_for_cost_injection",
+            return_value="gemini-2.5-pro",
+        ),
+        patch.object(
+            ProxyBaseLLMRequestProcessing,
+            "_process_chunk_with_cost_injection",
+            return_value=None,
+        ),
+    ):
+        received = []
+        async for chunk in PassThroughStreamingHandler.chunk_processor(
+            response=response,
+            request_body={"model": "gemini-2.5-pro"},
+            litellm_logging_obj=logging_obj,
+            endpoint_type=EndpointType.VERTEX_AI,
+            start_time=datetime.now(),
+            passthrough_success_handler_obj=MagicMock(),
+            url_route="/vertex_ai/v1/projects/p/locations/global/publishers/google/models/gemini-2.5-pro:streamRawPredict",
+        ):
+            received.append(chunk)
+        await asyncio.sleep(0)
+
+    assert received  # VERTEX_AI cost-injection branch was exercised
+    assert isinstance(logging_obj.completion_start_time, datetime)