scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 139

test(proxy): cover VERTEX_AI cost-injection branch in pass-through TTFT regression test

Browse source
This commit is contained in:
Bytechoreographer 2026-06-26 13:05:48 +08:00
parent cd37f2ed7f
commit 7818cd4aaf
1 changed files with 41 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_ttft.py
View file

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