From 983ab0f24eea1fdfe24d948e55909f9e09675cfb Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 19 Aug 2026 05:42:32 +0000
Subject: [PATCH] fix(tests): price the vertex batch cost test from rates it
 registers itself

The vertex batch aggregation test hardcoded gemini-3.6-flash's published
batch rates, so halving those rates in the model map left the assertion
expecting exactly twice the real cost and turned the misc unit-test job
red on litellm_internal_staging

It now registers its own vertex model whose batch rates are deliberately
not half its standard rates, so the assertion still proves batch pricing
is what gets applied while a future price refresh can no longer make it
stale
---
 .../test_litellm/batches/test_batch_utils.py  | 25 ++++++++++++++++---
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py
index 573882ebfca..453bdc10c70 100644
--- a/tests/test_litellm/batches/test_batch_utils.py
+++ b/tests/test_litellm/batches/test_batch_utils.py
@@ -19,6 +19,7 @@ import logging
 import os
 import sys
 from types import MappingProxyType
+from typing import Final
 
 import httpx
 import pytest
@@ -874,9 +875,25 @@ async def test_output_file_content_vertex_foreign_bucket_rejected_by_real_valida
 async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monkeypatch):
     import litellm.files.main as files_main
 
+    model: Final = "batch-priced-test-model"
+    litellm.register_model(
+        {
+            f"vertex_ai/{model}": {
+                "litellm_provider": "vertex_ai",
+                "mode": "chat",
+                "input_cost_per_token": 2e-06,
+                "output_cost_per_token": 8e-06,
+                "input_cost_per_token_batches": 3e-07,
+                "output_cost_per_token_batches": 5e-07,
+                "cache_read_input_token_cost": 0.0,
+                "cache_creation_input_token_cost": 0.0,
+            }
+        }
+    )
+
     rows = [
-        _vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5),
-        _vertex_openai_row("request-2", "gemini-3.6-flash", 20, 10),
+        _vertex_openai_row("request-1", model, 10, 5),
+        _vertex_openai_row("request-2", model, 20, 10),
     ]
 
     async def fake_afile_content(**kw):
@@ -891,9 +908,9 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
     )
 
     assert cost > 0
-    assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06)
+    assert cost == pytest.approx(30 * 3e-07 + 15 * 5e-07)
     assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
-    assert models == ["gemini-3.6-flash", "gemini-3.6-flash"]
+    assert models == [model, model]
 
 
 @pytest.mark.asyncio