From 49e0ddecdbca28786f791eab2b07f07024dde06f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 16 Jul 2025 08:55:07 -0700
Subject: [PATCH] test: update tests

---
 tests/llm_translation/test_gemini.py                   | 10 +++++-----
 .../llm_cost_calc/test_llm_cost_calc_utils.py          |  2 +-
 tests/test_litellm/test_cost_calculator.py             |  7 +------
 3 files changed, 7 insertions(+), 12 deletions(-)

diff --git a/tests/llm_translation/test_gemini.py b/tests/llm_translation/test_gemini.py
index c41d914e538..c4f796f65dc 100644
--- a/tests/llm_translation/test_gemini.py
+++ b/tests/llm_translation/test_gemini.py
@@ -23,7 +23,7 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest):
         return {"model": "gemini/gemini-2.0-flash"}
 
     def get_base_completion_call_args_with_reasoning_model(self) -> dict:
-        return {"model": "gemini/gemini-2.5-flash-preview-04-17"}
+        return {"model": "gemini/gemini-2.5-flash"}
 
     def test_tool_call_no_arguments(self, tool_call_no_arguments):
         """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""
@@ -34,7 +34,6 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest):
         result = convert_to_gemini_tool_call_invoke(tool_call_no_arguments)
         print(result)
 
-
     @pytest.mark.flaky(retries=3, delay=2)
     def test_url_context(self):
         from litellm.utils import supports_url_context
@@ -67,6 +66,7 @@ class TestGoogleAIStudioGemini(BaseLLMChatTest):
         print(f"response={response}")
 
 
+
 def test_gemini_context_caching_separate_messages():
     messages = [
         # System Message
@@ -151,13 +151,13 @@ def test_gemini_thinking():
     raw_request = return_raw_request(
         endpoint=CallTypes.completion,
         kwargs={
-            "model": "gemini/gemini-2.5-flash-preview-04-17",
+            "model": "gemini/gemini-2.5-flash",
             "messages": messages,
         },
     )
     assert reasoning_content in json.dumps(raw_request)
     response = completion(
-        model="gemini/gemini-2.5-flash-preview-04-17",
+        model="gemini/gemini-2.5-flash",
         messages=messages,  # make sure call works
     )
     print(response.choices[0].message)
@@ -173,7 +173,7 @@ def test_gemini_thinking_budget_0():
     raw_request = return_raw_request(
         endpoint=CallTypes.completion,
         kwargs={
-            "model": "gemini/gemini-2.5-flash-preview-04-17",
+            "model": "gemini/gemini-2.5-flash",
             "messages": [
                 {
                     "role": "user",
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 612662fd2c6..79707359c13 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
@@ -68,7 +68,7 @@ def test_reasoning_tokens_no_price_set():
 
 
 def test_reasoning_tokens_gemini():
-    model = "gemini-2.5-flash-preview-04-17"
+    model = "gemini-2.5-flash"
     custom_llm_provider = "gemini"
     os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
     litellm.model_cost = litellm.get_model_cost_map(url="")
diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index aa109f04c2d..18f78b90919 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -330,12 +330,7 @@ def test_cost_calculator_with_cache_creation():
 def test_bedrock_cost_calculator_comparison_with_without_cache():
     """Test that Bedrock caching reduces costs compared to non-cached requests"""
     from litellm import completion_cost
-    from litellm.types.utils import (
-        Choices,
-        Message,
-        PromptTokensDetailsWrapper,
-        Usage,
-    )
+    from litellm.types.utils import Choices, Message, PromptTokensDetailsWrapper, Usage
 
     # Response WITHOUT caching
     response_no_cache = ModelResponse(