test(azure_ai): cover the FLUX.2 billing fallbacks the PR adds

Pin what each fallback does so it can't silently change: a returned image
that isn't valid base64 bills the requested size, a custom-named FLUX.2
deployment with only a reference rate bills only its references, an
unlisted azure_ai deployment with no generated image price bills nothing
instead of raising, a non-ImageResponse is rejected, a FLUX.2 edit that
gets back a non-JSON body surfaces that body, a reference with no read()
is rejected, and get_image_dimensions falls back to the default size for
a header it can't read

tests/test_litellm/litellm_core_utils/test_token_counter.py

@@ -1667,6 +1667,7 @@ def test_image_dimensions_from_bytes_reads_each_header_format(image: bytes, expe
     [
         pytest.param(b"", id="empty"),
         pytest.param(b"BM" + b"\x00" * 30, id="unknown-format"),
+        pytest.param(_webp_bytes(b"ALPH", b"\x00" * 16), id="webp-without-an-image-chunk"),
         pytest.param(b"\x89PNG\r\n\x1a\n\x00\x00", id="png-truncated-before-ihdr"),
         pytest.param(b"\xff\xd8\xff\xe0\x00\x10JFIF", id="jpeg-truncated-inside-app0"),
         pytest.param(b"\xff\xd8\xff\xe0\x00\x04\x00\x00", id="jpeg-ends-before-sof"),
@@ -1708,3 +1709,17 @@ def test_image_dimensions_from_bytes_still_reads_a_jpeg_with_many_real_segments(
 def test_get_image_dimensions_still_raises_for_a_truncated_header(header: bytes) -> None:
     with pytest.raises((struct.error, TypeError)):
         get_image_dimensions(data="data:image/png;base64," + base64.b64encode(header).decode())
+
+
+@pytest.mark.parametrize(
+    "image",
+    [
+        pytest.param(b"BM" + b"\x00" * 30, id="unknown-format"),
+        pytest.param(b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1025 + _jpeg_sof(800, 600), id="pathological-jpeg"),
+    ],
+)
+def test_get_image_dimensions_falls_back_to_the_default_size_for_a_header_it_cannot_read(image: bytes) -> None:
+    assert get_image_dimensions(data="data:image/png;base64," + base64.b64encode(image).decode()) == (
+        litellm.constants.DEFAULT_IMAGE_WIDTH,
+        litellm.constants.DEFAULT_IMAGE_HEIGHT,
+    )

tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py

@@ -397,6 +397,18 @@ def test_flux2_image_edit_rejects_a_text_mode_upload_with_a_clear_error(tmp_path
         )
 
 
+def test_flux2_image_edit_rejects_a_reference_it_cannot_read():
+    with pytest.raises(ValueError, match="Unsupported image type"):
+        AzureFoundryFlux2ImageEditConfig().transform_image_edit_request(
+            model="FLUX.2-flex",
+            prompt="Make it a watercolor",
+            image="reference.png",
+            image_edit_optional_request_params={},
+            litellm_params={},
+            headers={},
+        )
+
+
 def test_flux2_image_edit_measures_a_stream_reference():
     uploaded: Final = io.BytesIO(_png(2048, 2048))
     response: Final = litellm.image_edit(
@@ -611,6 +623,23 @@ async def test_flux2_router_image_edit_bills_the_deployment_rates_with_a_logger_
     assert response._hidden_params["response_cost"] == pytest.approx(expected_cost)
 
 
+def test_flux2_image_edit_surfaces_a_response_that_is_not_json():
+    with pytest.raises(litellm.APIError, match="gateway timeout page"):
+        litellm.image_edit(
+            model="azure_ai/flux.2-pro",
+            image=_png(1024, 1024),
+            prompt="Make it a watercolor",
+            api_key="test-key",
+            api_base="https://example.services.ai.azure.com",
+            client=HTTPHandler(
+                client=httpx.Client(
+                    transport=httpx.MockTransport(lambda _request: httpx.Response(200, text="gateway timeout page"))
+                )
+            ),
+            size="1024x1024",
+        )
+
+
 def _edit_returning(image: bytes) -> Callable[[httpx.Request], httpx.Response]:
     def respond(request: httpx.Request) -> httpx.Response:
         return httpx.Response(200, json={"data": [{"b64_json": base64.b64encode(image).decode()}]})

tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py

@@ -13,6 +13,7 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
 from litellm.llms.azure.azure import AzureChatCompletion
 from litellm.llms.azure.image_generation import get_azure_image_generation_config
 from litellm.llms.azure.image_generation.http_utils import azure_deployment_image_generation_json_body
+from litellm.llms.azure_ai.image_generation.cost_calculator import cost_calculator as azure_ai_image_cost_calculator
 from litellm.llms.azure_ai.image_generation.flux_transformation import (
     AzureFoundryFluxImageGenerationConfig,
 )
@@ -427,6 +428,20 @@ def test_flux2_pro_bills_one_megapixel_for_a_size_it_cannot_measure(size: str):
     assert cost == pytest.approx(first)
 
 
+def test_flux2_pro_bills_the_requested_size_when_the_returned_image_is_not_valid_base64():
+    first, additional, _reference = _pro_megapixel_rates()
+
+    cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
+        model="flux.2-pro",
+        completion_response=ImageResponse(data=[ImageObject(b64_json="abc")]),
+        custom_llm_provider="azure_ai",
+        size="2048x1024",
+        call_type="image_generation",
+    )
+
+    assert cost == pytest.approx(first + additional)
+
+
 @pytest.mark.parametrize(("n", "billed_images"), ((2, 2), (None, 0)))
 def test_flux2_pro_bills_the_requested_image_count_when_the_response_lists_none(n: int | None, billed_images: int):
     first, _additional, _reference = _pro_megapixel_rates()
@@ -629,6 +644,24 @@ def test_custom_named_flux2_deployment_bills_its_own_megapixel_and_reference_rat
     assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2 + 2e-07 * 1024 * 1024 * 2)
 
 
+def test_custom_named_flux2_deployment_with_only_a_reference_rate_bills_only_its_references() -> None:
+    cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
+        model="my-flux2-prod",
+        completion_response=_edit_response((1024 * 1280,)),
+        custom_llm_provider="azure_ai",
+        size="1024x1280",
+        call_type="image_edit",
+        model_info={"input_cost_per_reference_pixel": 2e-07},
+    )
+
+    assert cost == pytest.approx(2e-07 * 1024 * 1024 * 2)
+
+
+def test_azure_ai_image_cost_calculator_rejects_a_response_that_is_not_an_image_response() -> None:
+    with pytest.raises(ValueError, match="must be of type ImageResponse"):
+        azure_ai_image_cost_calculator(model="flux.2-pro", image_response={"data": []})
+
+
 @pytest.mark.parametrize("model", ("FLUX-1.1-pro", "FLUX.1-Kontext-pro"))
 def test_flat_priced_flux_edit_ignores_reference_pixels(model: str):
     cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
@@ -672,6 +705,19 @@ def test_unlisted_azure_ai_model_bills_deployment_input_cost_per_pixel() -> None
     assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2)
 
 
+def test_unlisted_azure_ai_deployment_without_a_generated_image_price_bills_nothing() -> None:
+    cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
+        model="unlisted-flux-deployment",
+        completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]),
+        custom_llm_provider="azure_ai",
+        size="1024x1024",
+        call_type="image_generation",
+        model_info={"input_cost_per_reference_pixel": 1e-07},
+    )
+
+    assert cost == 0.0
+
+
 def test_flux2_response_preserves_mapped_dimensions():
     config = AzureFoundryFluxImageGenerationConfig()
     params = config.map_openai_params(