From ee2b51b2d08a7e4eee72b4929c8d929a92ae696a Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Wed, 7 Jan 2026 14:36:50 +0530
Subject: [PATCH] test_openai_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 68acb7ac7f..6914c7a37a 100644
--- a/tests/image_gen_tests/test_image_edits.py
+++ b/tests/image_gen_tests/test_image_edits.py
@@ -322,14 +322,23 @@ async def test_openai_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: