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 e8b3e117776..37330207bae 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
@@ -1456,12 +1456,12 @@ def test_character_count_billing_does_not_fill_prompt_token_gap():
     )
  
     prompt_cost, completion_cost = generic_cost_per_token(
-        model="gemini-2.0-pro",
+        model="gemini-2.0-flash-001",
         usage=usage,
         custom_llm_provider="vertex_ai",
     )
  
-    model_info = litellm.model_cost["gemini-2.0-pro"]
+    model_info = litellm.model_cost["gemini-2.0-flash-001"]
     expected_prompt_cost = (
         100 * model_info["input_cost_per_token"]
         + 1000 * model_info.get("input_cost_per_character", 0)
@@ -1498,12 +1498,12 @@ def test_video_length_billing_does_not_fill_prompt_token_gap():
     )
  
     prompt_cost, completion_cost = generic_cost_per_token(
-        model="gemini-2.0-pro",
+        model="gemini-2.0-flash-001",
         usage=usage,
         custom_llm_provider="vertex_ai",
     )
  
-    model_info = litellm.model_cost["gemini-2.0-pro"]
+    model_info = litellm.model_cost["gemini-2.0-flash-001"]
     expected_prompt_cost = (
         50 * model_info.get("input_cost_per_token", 0)
         + 12.0 * model_info.get("input_cost_per_video_per_second", 0)