diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index e7ce9a74797..c6b57fa7604 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -595,6 +595,8 @@ def test_completion_cost_image_generation_registered_deployment_price_keeps_map_
         deployment_id,
         {"mode": "image_generation", "litellm_provider": "gemini", "output_cost_per_image": 0.1},
     )
+    map_model: Final = "gemini/gemini-3.1-flash-image"
+    row: Final = litellm.model_cost[map_model]
     usage: Final = ImageUsage(
         input_tokens=10,
         input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=10),
@@ -604,7 +606,7 @@ def test_completion_cost_image_generation_registered_deployment_price_keeps_map_
 
     cost = completion_cost(
         completion_response=ImageResponse(data=[ImageObject(url="https://example.com/img.png")], usage=usage),
-        model="gemini/gemini-3.1-flash-image-preview",
+        model=map_model,
         custom_llm_provider="gemini",
         call_type="image_generation",
         custom_pricing=True,
@@ -612,7 +614,10 @@ def test_completion_cost_image_generation_registered_deployment_price_keeps_map_
         litellm_logging_obj=SimpleNamespace(litellm_params={"metadata": {"model_info": {"id": deployment_id}}}),
     )
 
-    assert cost == pytest.approx(10 * 5e-07 + 1290 * 6e-05)
+    expected: Final = (
+        usage.input_tokens * row["input_cost_per_token"] + usage.output_tokens * row["output_cost_per_image_token"]
+    )
+    assert cost == pytest.approx(expected)
 
 
 def test_completion_cost_image_generation_ignores_deployment_model_info_without_custom_pricing(