Address PR review nits: add tests, dedup MockResponse, env isolation

- Add test for static azure_ad_token string auth path (nit 1)
- Add test for azure_ad_token_provider exception propagation (nit 2)
- Add comment explaining api_key falsy guard behavior (nit 5)
- Extract MockResponse to module-level shared class (dedup 6 copies)
- Patch litellm.api_key, litellm.azure_key, and Azure env vars in AD
  token tests to ensure isolation from ambient CI/CD environment

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

litellm/llms/azure/image_edit/transformation.py

@@ -20,6 +20,8 @@ class AzureImageEditConfig(OpenAIImageEditConfig):
         api_base: Optional[str] = None,
     ) -> dict:
         _litellm_params = GenericLiteLLMParams(**(litellm_params or {}))
+        # Only override when api_key is truthy; empty string is not a valid key
+        # and _base_validate_azure_environment treats None as "no explicit key".
         if api_key:
             _litellm_params.api_key = api_key
         return BaseAzureLLM._base_validate_azure_environment(

tests/image_gen_tests/test_image_edits.py

@@ -33,6 +33,18 @@ class TestCustomLogger(CustomLogger):
         pass
 
 
+class MockResponse:
+    """Shared mock HTTP response used across image edit tests."""
+
+    def __init__(self, json_data, status_code):
+        self._json_data = json_data
+        self.status_code = status_code
+        self.text = json.dumps(json_data)
+
+    def json(self):
+        return self._json_data
+
+
 class BaseLLMImageEditTest(ABC):
     """
     Abstract base test class that enforces a common test across all image edit test classes.
@@ -256,15 +268,6 @@ async def test_azure_image_edit_litellm_sdk():
         ],
     }
 
-    class MockResponse:
-        def __init__(self, json_data, status_code):
-            self._json_data = json_data
-            self.status_code = status_code
-            self.text = json.dumps(json_data)
-
-        def json(self):
-            return self._json_data
-
     with patch(
         "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
         new_callable=AsyncMock,
@@ -374,15 +377,6 @@ async def test_openai_image_edit_cost_tracking():
         },
     }
 
-    class MockResponse:
-        def __init__(self, json_data, status_code):
-            self._json_data = json_data
-            self.status_code = status_code
-            self.text = json.dumps(json_data)
-
-        def json(self):
-            return self._json_data
-
     with patch(
         "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
         new_callable=AsyncMock,
@@ -464,15 +458,6 @@ async def test_azure_image_edit_cost_tracking():
         },
     }
 
-    class MockResponse:
-        def __init__(self, json_data, status_code):
-            self._json_data = json_data
-            self.status_code = status_code
-            self.text = json.dumps(json_data)
-
-        def json(self):
-            return self._json_data
-
     with patch(
         "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
         new_callable=AsyncMock,
@@ -743,15 +728,6 @@ async def test_image_edit_array_handling():
         ],
     }
 
-    class MockResponse:
-        def __init__(self, json_data, status_code):
-            self._json_data = json_data
-            self.status_code = status_code
-            self.text = json.dumps(json_data)
-
-        def json(self):
-            return self._json_data
-
     with patch(
         "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
         new_callable=AsyncMock,
@@ -796,22 +772,25 @@ async def test_azure_image_edit_azure_ad_token_auth():
         ],
     }
 
-    class MockResponse:
-        def __init__(self, json_data, status_code):
-            self._json_data = json_data
-            self.status_code = status_code
-            self.text = json.dumps(json_data)
-
-        def json(self):
-            return self._json_data
-
     fake_ad_token = "fake-azure-ad-token-12345"
     token_provider = lambda: fake_ad_token
 
-    with patch(
-        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
-        new_callable=AsyncMock,
-    ) as mock_post:
+    with (
+        patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            new_callable=AsyncMock,
+        ) as mock_post,
+        patch.object(litellm, "api_key", None),
+        patch.object(litellm, "azure_key", None),
+        patch.dict(
+            os.environ,
+            {},
+        ),
+    ):
+        # Remove ambient Azure env vars that would override AD token auth
+        os.environ.pop("AZURE_OPENAI_API_KEY", None)
+        os.environ.pop("AZURE_API_KEY", None)
+
         mock_post.return_value = MockResponse(mock_response, 200)
 
         test_api_base = "https://ai-api-gw-uae-north.openai.azure.com"
@@ -831,12 +810,16 @@ async def test_azure_image_edit_azure_ad_token_auth():
         headers = call_args.kwargs.get("headers", {})
 
         # Azure AD token auth should use Authorization: Bearer header
-        assert "Authorization" in headers, "Authorization header should be present for Azure AD auth"
-        assert headers["Authorization"] == f"Bearer {fake_ad_token}", (
-            f"Expected Bearer token with AD token, got {headers['Authorization']}"
-        )
+        assert (
+            "Authorization" in headers
+        ), "Authorization header should be present for Azure AD auth"
+        assert (
+            headers["Authorization"] == f"Bearer {fake_ad_token}"
+        ), f"Expected Bearer token with AD token, got {headers['Authorization']}"
         # api-key header should NOT be present
-        assert "api-key" not in headers, "api-key header should not be present for Azure AD auth"
+        assert (
+            "api-key" not in headers
+        ), "api-key header should not be present for Azure AD auth"
 
         ImageResponse.model_validate(result)
 
@@ -855,23 +838,26 @@ async def test_azure_image_edit_api_key_takes_precedence():
         ],
     }
 
-    class MockResponse:
-        def __init__(self, json_data, status_code):
-            self._json_data = json_data
-            self.status_code = status_code
-            self.text = json.dumps(json_data)
-
-        def json(self):
-            return self._json_data
-
     fake_ad_token = "fake-azure-ad-token-12345"
     token_provider = lambda: fake_ad_token
     test_api_key = "test-api-key-priority"
 
-    with patch(
-        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
-        new_callable=AsyncMock,
-    ) as mock_post:
+    with (
+        patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            new_callable=AsyncMock,
+        ) as mock_post,
+        patch.object(litellm, "api_key", None),
+        patch.object(litellm, "azure_key", None),
+        patch.dict(
+            os.environ,
+            {},
+        ),
+    ):
+        # Remove ambient Azure env vars so only the explicit api_key arg is used
+        os.environ.pop("AZURE_OPENAI_API_KEY", None)
+        os.environ.pop("AZURE_API_KEY", None)
+
         mock_post.return_value = MockResponse(mock_response, 200)
 
         result = await aimage_edit(
@@ -889,11 +875,116 @@ async def test_azure_image_edit_api_key_takes_precedence():
         headers = call_args.kwargs.get("headers", {})
 
         # API key should take precedence — uses api-key header
-        assert "api-key" in headers, "api-key header should be present when API key is provided"
+        assert (
+            "api-key" in headers
+        ), "api-key header should be present when API key is provided"
         assert headers["api-key"] == test_api_key
         # Authorization header should NOT be present
-        assert "Authorization" not in headers, (
-            "Authorization header should not be set when API key is available"
-        )
+        assert (
+            "Authorization" not in headers
+        ), "Authorization header should not be set when API key is available"
 
         ImageResponse.model_validate(result)
+
+
+@pytest.mark.asyncio
+async def test_azure_image_edit_azure_ad_token_string_auth():
+    """Test Azure image edit uses a static azure_ad_token string when no API key is provided."""
+    from litellm import aimage_edit
+
+    mock_response = {
+        "created": 1589478378,
+        "data": [
+            {
+                "b64_json": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg=="
+            }
+        ],
+    }
+
+    static_ad_token = "static-azure-ad-token-67890"
+
+    with (
+        patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            new_callable=AsyncMock,
+        ) as mock_post,
+        patch.object(litellm, "api_key", None),
+        patch.object(litellm, "azure_key", None),
+        patch.dict(
+            os.environ,
+            {},
+        ),
+    ):
+        os.environ.pop("AZURE_OPENAI_API_KEY", None)
+        os.environ.pop("AZURE_API_KEY", None)
+
+        mock_post.return_value = MockResponse(mock_response, 200)
+
+        result = await aimage_edit(
+            prompt="Edit this image",
+            model="azure/gpt-image-1",
+            api_base="https://test.openai.azure.com",
+            api_version="2025-04-01-preview",
+            azure_ad_token=static_ad_token,
+            image=TEST_IMAGES,
+        )
+
+        mock_post.assert_called_once()
+        call_args = mock_post.call_args
+        headers = call_args.kwargs.get("headers", {})
+
+        # Static AD token should use Authorization: Bearer header
+        assert (
+            "Authorization" in headers
+        ), "Authorization header should be present for static azure_ad_token"
+        assert (
+            headers["Authorization"] == f"Bearer {static_ad_token}"
+        ), f"Expected Bearer token with static AD token, got {headers['Authorization']}"
+        # api-key header should NOT be present
+        assert (
+            "api-key" not in headers
+        ), "api-key header should not be present for static azure_ad_token auth"
+
+        ImageResponse.model_validate(result)
+
+
+@pytest.mark.asyncio
+async def test_azure_image_edit_ad_token_provider_raises():
+    """Test that an exception from azure_ad_token_provider propagates to the caller."""
+    from litellm import aimage_edit
+
+    def failing_provider():
+        raise ValueError("Token refresh failed")
+
+    with (
+        patch(
+            "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+            new_callable=AsyncMock,
+        ) as mock_post,
+        patch.object(litellm, "api_key", None),
+        patch.object(litellm, "azure_key", None),
+        patch.dict(
+            os.environ,
+            {},
+        ),
+    ):
+        os.environ.pop("AZURE_OPENAI_API_KEY", None)
+        os.environ.pop("AZURE_API_KEY", None)
+
+        mock_post.return_value = MockResponse({"error": "should not reach"}, 200)
+
+        with pytest.raises(
+            litellm.exceptions.APIConnectionError,
+            match="Failed to get Azure AD token",
+        ):
+            await aimage_edit(
+                prompt="Edit this image",
+                model="azure/gpt-image-1",
+                api_base="https://test.openai.azure.com",
+                api_version="2025-04-01-preview",
+                azure_ad_token_provider=failing_provider,
+                image=TEST_IMAGES,
+            )
+
+        # The HTTP call should never have been made
+        mock_post.assert_not_called()