From ca537f5a7bc0abfe566f649d22a4214de25ee7ea Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 7 Jan 2026 17:58:22 +0530 Subject: [PATCH] TestAzureAIFlux2ImageEdit --- tests/image_gen_tests/test_image_edits.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/image_gen_tests/test_image_edits.py b/tests/image_gen_tests/test_image_edits.py index 62f8341e0a4..810bd80a5b0 100644 --- a/tests/image_gen_tests/test_image_edits.py +++ b/tests/image_gen_tests/test_image_edits.py @@ -143,6 +143,23 @@ class TestOpenAIImageEditDallE2(BaseLLMImageEditTest): } +class TestAzureAIFlux2ImageEdit(BaseLLMImageEditTest): + """ + Concrete implementation of BaseLLMImageEditTest for Azure AI FLUX 2 image edits. + FLUX 2 uses JSON with base64 image instead of multipart/form-data. + """ + + def get_base_image_edit_call_args(self) -> dict: + """Return base call args for Azure AI FLUX 2 image edit""" + return { + "model": "azure_ai/flux.2-pro", + "image": SINGLE_TEST_IMAGE, + "api_base": os.getenv("AZURE_AI_API_BASE", "https://litellm-ci-cd-prod.services.ai.azure.com"), + "api_key": os.getenv("AZURE_AI_API_KEY"), + "api_version": "preview", + } + + @pytest.mark.flaky(retries=3, delay=2) @pytest.mark.asyncio async def test_openai_image_edit_litellm_router():