From 1b72cdcfc0ae9051e34edb06ae48f4828bdc0a5e Mon Sep 17 00:00:00 2001
From: kerry <kerry@berri.ai>
Date: Tue, 22 Sep 2026 20:45:29 +0000
Subject: [PATCH] test(integration): azure_ai FLUX.2-flex image generation
 targets the flex provider path with the BFL body (Pylon #7092)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 tests/integration/contracts.json              |  3 ++
 .../test_azure_ai_flux2_image_wire.py         | 48 +++++++++++++++++++
 2 files changed, 51 insertions(+)
 create mode 100644 tests/integration/providers/test_azure_ai_flux2_image_wire.py

diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json
index f7b4db7ad9d..e49bba7e90c 100644
--- a/tests/integration/contracts.json
+++ b/tests/integration/contracts.json
@@ -226,6 +226,9 @@
     "tests/integration/providers/test_fal_ai_image_wire.py::test_fal_flux_dev_generation_without_deployment_api_base_uses_global_api_base": [
       "other.provider_wire.fal_ai.global_api_base_routes_image_generation"
     ],
+    "tests/integration/providers/test_azure_ai_flux2_image_wire.py::test_azure_flux2_flex_generation_hits_flex_provider_path_not_pro": [
+      "other.provider_wire.azure_ai.flux2_flex_generation_targets_flex_path_with_bfl_body"
+    ],
     "tests/integration/providers/test_xiaomi_mimo_wire.py::test_xiaomi_mimo_nonstream_surfaces_reasoning_and_charges_registry_price[mimo-v2.6-pro]": [
       "other.provider_wire.xiaomi_mimo.reasoning_content_and_registry_pricing"
     ],
diff --git a/tests/integration/providers/test_azure_ai_flux2_image_wire.py b/tests/integration/providers/test_azure_ai_flux2_image_wire.py
new file mode 100644
index 00000000000..59f125463a6
--- /dev/null
+++ b/tests/integration/providers/test_azure_ai_flux2_image_wire.py
@@ -0,0 +1,48 @@
+import json
+from typing import Final
+
+import pytest
+from integration._support.client import Gateway
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue, TypeAdapter
+
+_FLEX_MODEL: Final = "azure_ai/FLUX.2-flex"
+_PROMPT: Final = "a red fox in the snow"
+_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+
+
+@pytest.mark.covers("other.provider_wire.azure_ai.flux2_flex_generation_targets_flex_path_with_bfl_body")
+def test_azure_flux2_flex_generation_hits_flex_provider_path_not_pro(gateway: Gateway) -> None:
+    def respond(request: Request) -> Reply:
+        assert request.method == "POST"
+        assert request.target == "/providers/blackforestlabs/v1/flux-2-flex?api-version=preview"
+        assert request.headers["api-key"] == "synthetic-azure-key"
+        assert _JSON_OBJECT.validate_json(request.body) == {
+            "model": "FLUX.2-flex",
+            "prompt": _PROMPT,
+            "num_images": 2,
+            "width": 1536,
+            "height": 1024,
+            "guidance": 4.5,
+            "steps": 32,
+        }
+        return Reply(body=json.dumps({"data": [{"b64_json": "aW1n"}, {"b64_json": "aW1n"}]}).encode())
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model=_FLEX_MODEL, api_base=wire.url, api_key="synthetic-azure-key", api_version="preview"
+        )
+        response: Final = gateway.request(
+            "POST",
+            "/v1/images/generations",
+            {"model": model, "prompt": _PROMPT, "n": 2, "size": "1536x1024", "guidance": 4.5, "steps": 32},
+        )
+        assert response.status_code == 200, response.text
+        payload: Final = _JSON_OBJECT.validate_json(response.content)
+        assert payload["data"] == [
+            {"url": None, "b64_json": "aW1n", "revised_prompt": None, "provider_specific_fields": None},
+            {"url": None, "b64_json": "aW1n", "revised_prompt": None, "provider_specific_fields": None},
+        ]
+        assert [(request.method, request.target) for request in wire.drain()] == [
+            ("POST", "/providers/blackforestlabs/v1/flux-2-flex?api-version=preview")
+        ]