From 13691426c4610f10da9627388a53419196a989ee Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:52:05 -0700 Subject: [PATCH] test(cost_calculator): point image-generation deployment price test at a live gemini row (#42615) The test priced gemini/gemini-3.1-flash-image-preview, which #42435 removed from the cost map as deprecated, so the calculator had no per-token rates to keep and the hardcoded expected value no longer matched. Price the live gemini/gemini-3.1-flash-image row instead and derive the expected cost from that row in litellm.model_cost, so a rate change on it cannot break the test while a calculator that drops the map's token rates still fails it. Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- tests/test_litellm/test_cost_calculator.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index e7ce9a74797..c6b57fa7604 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -595,6 +595,8 @@ def test_completion_cost_image_generation_registered_deployment_price_keeps_map_ deployment_id, {"mode": "image_generation", "litellm_provider": "gemini", "output_cost_per_image": 0.1}, ) + map_model: Final = "gemini/gemini-3.1-flash-image" + row: Final = litellm.model_cost[map_model] usage: Final = ImageUsage( input_tokens=10, input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=10), @@ -604,7 +606,7 @@ def test_completion_cost_image_generation_registered_deployment_price_keeps_map_ cost = completion_cost( completion_response=ImageResponse(data=[ImageObject(url="https://example.com/img.png")], usage=usage), - model="gemini/gemini-3.1-flash-image-preview", + model=map_model, custom_llm_provider="gemini", call_type="image_generation", custom_pricing=True, @@ -612,7 +614,10 @@ def test_completion_cost_image_generation_registered_deployment_price_keeps_map_ litellm_logging_obj=SimpleNamespace(litellm_params={"metadata": {"model_info": {"id": deployment_id}}}), ) - assert cost == pytest.approx(10 * 5e-07 + 1290 * 6e-05) + expected: Final = ( + usage.input_tokens * row["input_cost_per_token"] + usage.output_tokens * row["output_cost_per_image_token"] + ) + assert cost == pytest.approx(expected) def test_completion_cost_image_generation_ignores_deployment_model_info_without_custom_pricing(