diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index 6f9abdf7ff6..64a21a35097 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/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( diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 1fc8da88772..2ba44efcaff 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/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()