From ad899986b9ac7ba62e9376ee2e66b34b74b274ab Mon Sep 17 00:00:00 2001
From: Praveen11558 <44603409+Praveen11558@users.noreply.github.com>
Date: Mon, 23 Mar 2026 16:36:23 +0530
Subject: [PATCH] Update test_llm_cost_calc_utils.py

---
 .../llm_cost_calc/test_llm_cost_calc_utils.py                 | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
index 61679a6fa21..dcedd3d7df3 100644
--- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
+++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
@@ -1505,10 +1505,10 @@ def test_video_length_billing_does_not_fill_prompt_token_gap():
  
     model_info = litellm.model_cost["gemini-2.5-pro"]
     expected_prompt_cost = (
-        50 * model_info["input_cost_per_token"]
+        50 * model_info.get("input_cost_per_token", 0)
         + 12.0 * model_info.get("input_cost_per_video_per_second", 0)
     )
-    expected_completion_cost = 10 * model_info["output_cost_per_token"]
+    expected_completion_cost = 10 * model_info.get("output_cost_per_token", 0)
  
     assert prompt_cost == pytest.approx(expected_prompt_cost), (
         f"Expected prompt_cost={expected_prompt_cost}, got {prompt_cost}. "