From 15e8a35feb0bbdbe113c2389bf00c0738cb8bbc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 17:54:44 -0700 Subject: [PATCH] test: separate gemini entries storing the promo rate from those storing list The 20% promotion that runs to 2027-01-31 covers every gemini model, not just the 2.5 pair, so a constant naming two of them implied the other four were exempt. Six covered entries live in the registry: two store the discounted rate and four store list, which is a pre-existing overcharge this branch does not touch, since it only adds cache fields and derives them from each entry's own input rate. Name both groups for what they store, pin the expiry, and tighten the tolerance to 2e-4. --- .../test_databricks_cost_calculator.py | 40 +++++++++++++++---- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index d5abbc547e5..21f047b753c 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -63,10 +63,17 @@ PUBLISHED_DBU_PER_MILLION: Final = { "databricks/databricks-gemini-2-5-flash": ("5.357", "44.643", "5.357", "0.536"), } PROMOTIONAL_DISCOUNT: Final = 0.80 -PROMOTIONALLY_DISCOUNTED_MODELS: Final = ( +PROMOTION_EXPIRES: Final = "2027-01-31" +ENTRIES_STORING_PROMOTIONAL_RATE: Final = ( "databricks/databricks-gemini-2-5-pro", "databricks/databricks-gemini-2-5-flash", ) +ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION: Final = ( + "databricks/databricks-gemini-3-1-pro", + "databricks/databricks-gemini-3-pro", + "databricks/databricks-gemini-3-flash", + "databricks/databricks-gemini-3-1-flash-lite", +) CACHE_FIELDS: Final = ("cache_creation_input_token_cost", "cache_read_input_token_cost") @@ -136,7 +143,7 @@ def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, mod assert info[field] == _dollars_per_token(dbu_per_million), field -@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(PROMOTIONALLY_DISCOUNTED_MODELS))) +@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(ENTRIES_STORING_PROMOTIONAL_RATE))) def test_cache_rates_derive_from_published_cache_dbu(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) cache_dbu_per_million: Final = PUBLISHED_DBU_PER_MILLION[model][2:] @@ -222,17 +229,36 @@ def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: No assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field -@pytest.mark.parametrize("model", PROMOTIONALLY_DISCOUNTED_MODELS) -def test_promotionally_discounted_entries_price_below_the_published_table( +@pytest.mark.parametrize("model", ENTRIES_STORING_PROMOTIONAL_RATE) +def test_entries_storing_the_promotional_rate_price_below_the_published_table( local_model_cost_map: None, model: str, ) -> None: info: Final = _model_info(model) input_dbu, output_dbu, _, _ = PUBLISHED_DBU_PER_MILLION[model] + expiry_hint: Final = f"the gemini promotion expires {PROMOTION_EXPIRES}, after which the list rate applies" - assert info["input_cost_per_token"] == pytest.approx(_dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=1e-3) + assert info["input_cost_per_token"] == pytest.approx( + _dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 + ), expiry_hint assert info["output_cost_per_token"] == pytest.approx( - _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=1e-3 - ) + _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 + ), expiry_hint assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"]) + + +@pytest.mark.parametrize("model", ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION) +def test_entries_storing_the_list_rate_bill_above_the_promotional_price( + local_model_cost_map: None, + model: str, +) -> None: + info: Final = _model_info(model) + input_dbu, _, _, _ = PUBLISHED_DBU_PER_MILLION[model] + list_rate: Final = _dollars_per_token(input_dbu) + + assert info["input_cost_per_token"] == pytest.approx(list_rate, rel=2e-4), ( + f"{model} moved off the list rate; if it now stores the discount that runs to " + f"{PROMOTION_EXPIRES}, move it into ENTRIES_STORING_PROMOTIONAL_RATE" + ) + assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"])