From 945b10ded467e53fc3c9b8df0329dbc55591a56e Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Tue, 12 May 2026 16:16:59 -0700
Subject: [PATCH] fix(tests): drop dall-e-only test classes; route live image
 tests via gpt-image-1

Second wave of failures from the 2026-05-12 DALL-E shutdown:
- tests/image_gen_tests/test_image_edits.py::TestOpenAIImageEditDallE2
  and tests/image_gen_tests/test_image_generation.py::TestOpenAIDalle3
  are explicitly named for the deprecated models and can't pass; remove.
  gpt-image-1 coverage already exists in sibling classes.
- tests/local_testing/test_router.py image gen tests use dall-e-3 only
  as a routing example; swap to gpt-image-1.
- tests/local_testing/test_custom_callback_input.py image_generation
  success/failure paths swapped to gpt-image-1.
---
 tests/image_gen_tests/test_image_edits.py         | 14 --------------
 tests/image_gen_tests/test_image_generation.py    |  5 -----
 tests/local_testing/test_custom_callback_input.py |  4 ++--
 tests/local_testing/test_router.py                | 12 ++++++------
 4 files changed, 8 insertions(+), 27 deletions(-)

diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py
index dcb04c597e..6900bacdb9 100644
--- a/tests/image_gen_tests/test_image_edits.py
+++ b/tests/image_gen_tests/test_image_edits.py
@@ -133,20 +133,6 @@ class TestOpenAIImageEditGPTImage1(BaseLLMImageEditTest):
         }
 
 
-class TestOpenAIImageEditDallE2(BaseLLMImageEditTest):
-    """
-    Concrete implementation of BaseLLMImageEditTest for OpenAI DALL-E-2 image edits.
-    DALL-E-2 only supports a single image (not an array).
-    """
-
-    def get_base_image_edit_call_args(self) -> dict:
-        """Return base call args for OpenAI DALL-E-2 image edit (single image only)"""
-        return {
-            "model": "dall-e-2",
-            "image": SINGLE_TEST_IMAGE,
-        }
-
-
 class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest):
     """
     Concrete implementation of BaseLLMImageEditTest for Azure AI FLUX 2 image edits.
diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py
index 5152e3e012..873777189c 100644
--- a/tests/image_gen_tests/test_image_generation.py
+++ b/tests/image_gen_tests/test_image_generation.py
@@ -163,11 +163,6 @@ class TestBedrockNovaCanvasColorGuidedGeneration(BaseImageGenTest):
         }
 
 
-class TestOpenAIDalle3(BaseImageGenTest):
-    def get_base_image_generation_call_args(self) -> dict:
-        return {"model": "dall-e-3"}
-
-
 class TestOpenAIGPTImage1(BaseImageGenTest):
     def get_base_image_generation_call_args(self) -> dict:
         return {"model": "gpt-image-1"}
diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py
index 15a2975bec..545039e60b 100644
--- a/tests/local_testing/test_custom_callback_input.py
+++ b/tests/local_testing/test_custom_callback_input.py
@@ -930,7 +930,7 @@ def test_image_generation_openai():
 
         response = litellm.image_generation(
             prompt="A cute baby sea otter",
-            model="openai/dall-e-3",
+            model="openai/gpt-image-1",
             api_key=os.getenv("OPENAI_API_KEY"),
         )
 
@@ -948,7 +948,7 @@ def test_image_generation_openai():
         try:
             response = litellm.image_generation(
                 prompt="A cute baby sea otter",
-                model="dall-e-2",
+                model="gpt-image-1",
                 api_key="my-bad-api-key",
             )
         except Exception:
diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py
index f7885fb8a0..d6b239c79c 100644
--- a/tests/local_testing/test_router.py
+++ b/tests/local_testing/test_router.py
@@ -995,15 +995,15 @@ async def test_aimg_gen_on_router():
     try:
         model_list = [
             {
-                "model_name": "dall-e-3",
+                "model_name": "gpt-image-1",
                 "litellm_params": {
-                    "model": "dall-e-3",
+                    "model": "gpt-image-1",
                 },
             }
         ]
         router = Router(model_list=model_list, num_retries=3)
         response = await router.aimage_generation(
-            model="dall-e-3", prompt="A cute baby sea otter"
+            model="gpt-image-1", prompt="A cute baby sea otter"
         )
         print(response)
         assert len(response.data) > 0
@@ -1030,15 +1030,15 @@ def test_img_gen_on_router():
     try:
         model_list = [
             {
-                "model_name": "dall-e-3",
+                "model_name": "gpt-image-1",
                 "litellm_params": {
-                    "model": "dall-e-3",
+                    "model": "gpt-image-1",
                 },
             }
         ]
         router = Router(model_list=model_list)
         response = router.image_generation(
-            model="dall-e-3", prompt="A cute baby sea otter"
+            model="gpt-image-1", prompt="A cute baby sea otter"
         )
         print(response)
         assert len(response.data) > 0