From 2d2b7e8fa0488a5015c1589b003f764443670883 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Tue, 22 Sep 2026 18:37:41 -0700
Subject: [PATCH] test(cost): repoint the Azure image cost test at gpt-image-2
 (#42631)

#42435 removed the OpenAI dall-e-3 cost-map row that
test_dalle_3_azure_cost_tracking pinned, and Azure retired DALL-E 3
inference on 2026-03-04, so the test now mirrors a real Azure
gpt-image-2 response and derives the expected cost from the cost map.

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
---
 tests/local_testing/test_completion_cost.py | 63 ++++++++++++---------
 1 file changed, 35 insertions(+), 28 deletions(-)

diff --git a/tests/local_testing/test_completion_cost.py b/tests/local_testing/test_completion_cost.py
index 3ce99f893d8..c24e6c32369 100644
--- a/tests/local_testing/test_completion_cost.py
+++ b/tests/local_testing/test_completion_cost.py
@@ -5,7 +5,7 @@ import litellm.cost_calculator
 
 import asyncio
 import time
-from typing import Optional
+from typing import Final, Optional
 from unittest.mock import MagicMock, patch
 import pytest
 
@@ -21,6 +21,9 @@ import json
 import httpx
 from litellm.types.utils import PromptTokensDetails
 from litellm.litellm_core_utils.litellm_logging import CustomLogger
+from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
+    convert_to_model_response_object,
+)
 
 
 class CustomLoggingHandler(CustomLogger):
@@ -328,35 +331,39 @@ def test_whisper_azure():
     assert round(cost, 5) == round(expected_cost, 5)
 
 
-def test_dalle_3_azure_cost_tracking():
-    litellm.set_verbose = True
-    # model = "azure/dall-e-3-test"
-    # response = litellm.image_generation(
-    #     model=model,
-    #     prompt="A cute baby sea otter",
-    #     api_version="2023-12-01-preview",
-    #     api_base=os.getenv("AZURE_SWEDEN_API_BASE"),
-    #     api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
-    #     base_model="dall-e-3",
-    # )
-    # print(f"response: {response}")
-    response = litellm.ImageResponse(
-        created=1710265780,
-        data=[
-            {
-                "b64_json": None,
-                "revised_prompt": "A close-up image of an adorable baby sea otter. Its fur is thick and fluffy to provide buoyancy and insulation against the cold water. Its eyes are round, curious and full of life. It's lying on its back, floating effortlessly on the calm sea surface under the warm sun. Surrounding the otter are patches of colorful kelp drifting along the gentle waves, giving the scene a touch of vibrancy. The sea otter has its small paws folded on its chest, and it seems to be taking a break from its play.",
-                "url": "test-azure-blob-url-with-sas-token",
-            }
-        ],
+def test_gpt_image_2_azure_cost_tracking():
+    azure_image_generation_response: Final = {
+        "created": 1758585600,
+        "data": [{"b64_json": "iVBORw0KGgo=", "revised_prompt": None, "url": None}],
+        "output_format": "png",
+        "quality": "low",
+        "size": "1024x1024",
+        "usage": {
+            "input_tokens": 12,
+            "input_tokens_details": {"image_tokens": 0, "text_tokens": 12},
+            "output_tokens": 196,
+            "output_tokens_details": {"image_tokens": 196, "text_tokens": 0},
+            "total_tokens": 208,
+        },
+    }
+    response: Final = convert_to_model_response_object(
+        response_object=azure_image_generation_response,
+        model_response_object=litellm.ImageResponse(),
+        response_type="image_generation",
+        hidden_params={"model": "gpt-image-2", "custom_llm_provider": "azure"},
     )
-    response.usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
-    response._hidden_params = {"model": "dall-e-3", "model_id": None}
-    print(f"response hidden params: {response._hidden_params}")
-    cost = litellm.completion_cost(
-        completion_response=response, call_type="image_generation"
+
+    cost: Final = litellm.completion_cost(
+        completion_response=response,
+        model="azure/my-gpt-image-2-deployment",
+        custom_llm_provider="azure",
+        base_model="gpt-image-2",
+        call_type="image_generation",
     )
-    assert cost > 0
+
+    pricing: Final = litellm.model_cost["azure/gpt-image-2"]
+    expected_cost: Final = pricing["input_cost_per_token"] * 12 + pricing["output_cost_per_image_token"] * 196
+    assert round(cost, 8) == round(expected_cost, 8)
 
 
 def test_replicate_llama3_cost_tracking():