test(streaming): use a fixed start time and check the estimated cost

This commit is contained in:
Zuraiz Anjum 2026-10-04 16:17:08 +05:00
parent 9e51552d60
commit d80daa4827

View file

@ -1,3 +1,4 @@
import datetime
import json
import time
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@ -2378,7 +2379,7 @@ async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
messages=messages,
stream=True,
call_type="completion",
start_time=time.time(),
start_time=datetime.datetime(2025, 3, 3, 21, 38, 10),
litellm_call_id="no-usage-call",
function_id="no-usage-fn",
),
@ -2397,6 +2398,10 @@ async def test_stream_without_provider_usage_falls_back_to_the_token_estimate(
hidden_usage: Final = collected[-1]._hidden_params["usage"]
assert (hidden_usage.prompt_tokens, hidden_usage.completion_tokens) == expected
rates: Final = litellm.model_cost[model]
spend: Final = expected[0] * rates["input_cost_per_token"] + expected[1] * rates["output_cost_per_token"]
assert litellm.completion_cost(completion_response=assembled) == pytest.approx(spend)
@pytest.mark.asyncio
async def test_custom_stream_wrapper_aclose():