From c3341a1e187df9c56bf008703a8e866c72971ff5 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 2 Apr 2025 20:56:20 -0700
Subject: [PATCH] test fixes - azure deprecated dall-e-2

---
 .../test_custom_callback_input.py              |  8 ++++----
 tests/local_testing/test_router.py             | 18 +-----------------
 2 files changed, 5 insertions(+), 21 deletions(-)

diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py
index b0ebcf7767..222572935b 100644
--- a/tests/local_testing/test_custom_callback_input.py
+++ b/tests/local_testing/test_custom_callback_input.py
@@ -1133,10 +1133,10 @@ def test_image_generation_openai():
 
         response = litellm.image_generation(
             prompt="A cute baby sea otter",
-            model="azure/",
-            api_base=os.getenv("AZURE_API_BASE"),
-            api_key=os.getenv("AZURE_API_KEY"),
-            api_version="2023-06-01-preview",
+            model="azure/dall-e-3-test",
+            api_version="2023-12-01-preview",
+            api_base=os.getenv("AZURE_SWEDEN_API_BASE"),
+            api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
         )
 
         print(f"response: {response}")
diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py
index 20a2f28c95..68a79f94a6 100644
--- a/tests/local_testing/test_router.py
+++ b/tests/local_testing/test_router.py
@@ -1136,16 +1136,7 @@ async def test_aimg_gen_on_router():
                     "api_base": os.getenv("AZURE_SWEDEN_API_BASE"),
                     "api_key": os.getenv("AZURE_SWEDEN_API_KEY"),
                 },
-            },
-            {
-                "model_name": "dall-e-2",
-                "litellm_params": {
-                    "model": "azure/",
-                    "api_version": "2023-06-01-preview",
-                    "api_base": os.getenv("AZURE_API_BASE"),
-                    "api_key": os.getenv("AZURE_API_KEY"),
-                },
-            },
+            }
         ]
         router = Router(model_list=model_list, num_retries=3)
         response = await router.aimage_generation(
@@ -1153,13 +1144,6 @@ async def test_aimg_gen_on_router():
         )
         print(response)
         assert len(response.data) > 0
-
-        response = await router.aimage_generation(
-            model="dall-e-2", prompt="A cute baby sea otter"
-        )
-        print(response)
-        assert len(response.data) > 0
-
         router.reset()
     except litellm.InternalServerError as e:
         pass