From 49343c421badd248ff03458d0d1332d45751afb3 Mon Sep 17 00:00:00 2001 From: Shreshth Kharbanda Date: Sat, 26 Sep 2026 00:33:15 +0000 Subject: [PATCH] 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 --- .../litellm_core_utils/test_token_counter.py | 15 ++++++ ...test_azure_ai_image_edit_transformation.py | 29 ++++++++++++ .../test_azure_ai_flux2_image_generation.py | 46 +++++++++++++++++++ 3 files changed, 90 insertions(+) diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 9f1f13bc5a8..cfa5b00e6ca 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/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, + ) diff --git a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py b/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py index f5083094c65..4cfcb83918b 100644 --- a/tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py +++ b/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()}]}) diff --git a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py b/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py index 16f94ee6fd3..2742fcbab4a 100644 --- a/tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py +++ b/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(