From bb1f65c546b8a91099ba4e586c0325e0eb06475b Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 7 Jan 2026 15:22:13 +0530
Subject: [PATCH] fix: test_azure_image_edit_cost_tracking

---
 tests/image_gen_tests/test_image_edits.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py
index 6914c7a37a..62f8341e0a 100644
--- a/tests/image_gen_tests/test_image_edits.py
+++ b/tests/image_gen_tests/test_image_edits.py
@@ -410,14 +410,23 @@ async def test_azure_image_edit_cost_tracking():
     litellm.logging_callback_manager._reset_all_callbacks()
     litellm.callbacks = [test_custom_logger]
     
-    # Mock response for Azure image edit
+    # Mock response for Azure image edit with usage data for cost tracking
     mock_response = {
         "created": 1589478378,
         "data": [
             {
                 "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
             }
-        ]
+        ],
+        "usage": {
+            "total_tokens": 1100,
+            "input_tokens": 100,
+            "input_tokens_details": {
+                "image_tokens": 50,
+                "text_tokens": 50
+            },
+            "output_tokens": 1000
+        }
     }
 
     class MockResponse: