scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

test: close the FLUX.2 billing and JPEG header gaps mutation testing found

Browse source 
A mutation run over every changed line showed several behaviors nothing
pinned. This adds cases for an OpenAI deployment's own edit price reaching the
proxy's cost logging, every JPEG start-of-frame marker and its look-alikes,
per-segment skipping, and the segment and fill-byte limits at their
boundaries. It also covers a landscape lone reference, a 1-pixel reference,
unusable mapped dimensions, the -x- size spelling, exact pixel pricing for
non-FLUX.2 models, and the warnings logged when a reference is unmeasured or
its counts are malformed

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01G136StYRutjpSqCjgLF4q9
This commit is contained in:
Claude 2026-09-26 01:22:37 +00:00
parent add9b0b9ea
commit cf3225ef46
No known key found for this signature in database
5 changed files with 190 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
tests/unit/images/test_main.py Normal file
View file

@ -0,0 +1,55 @@
from datetime import datetime
from typing import Final
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx import llm_http_handler as llm_http_handler_module
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
PNG_BYTES: Final = b"\x89PNG\r\n\x1a\nfakepng"
def _edit_ok(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]})
@pytest.mark.asyncio
@pytest.mark.parametrize("model", ("openai/dall-e-2", "openai/unpriced-image-model"))
async def test_router_image_edit_bills_the_deployment_price_with_a_logger_built_before_routing(
monkeypatch: pytest.MonkeyPatch, model: str
):
client: Final = AsyncHTTPHandler()
client.client = httpx.AsyncClient(transport=httpx.MockTransport(_edit_ok))
monkeypatch.setattr(llm_http_handler_module, "get_async_httpx_client", lambda **_kwargs: client)
router: Final = litellm.Router(
model_list=[
{
"model_name": "priced-edit-deployment",
"litellm_params": {
"model": model,
"api_base": "https://edit.example/v1",
"api_key": "sk-test",
"output_cost_per_image": 0.5,
},
}
]
)
request: Final = {
"model": "priced-edit-deployment",
"prompt": "add a hat",
"image": PNG_BYTES,
"size": "1024x1024",
"litellm_call_id": "proxy-call-id",
}
logging_obj, routed_request = litellm.utils.function_setup(
original_function="aimage_edit",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**request,
)
response: Final = await router.aimage_edit(**routed_request, litellm_logging_obj=logging_obj)
assert response._hidden_params["response_cost"] == pytest.approx(0.5)

84
tests/unit/litellm_core_utils/test_token_counter.py
View file

@ -24,6 +24,8 @@ import litellm.constants
from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.token_counter import (
MAX_JPEG_FILL_BYTES,
MAX_JPEG_HEADER_SEGMENTS,
_get_exact_count_function,
_get_extrapolating_count_function,
_get_tiktoken_count_function,
@ -1517,26 +1519,77 @@ def test_image_dimensions_from_bytes_returns_none_for_unreadable_headers(image:
assert image_dimensions_from_bytes(image) is None
def _jpeg_sof(width: int, height: int) -> bytes:
return b"\xff\xc0" + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01"
def _jpeg_sof(width: int, height: int, marker: bytes = b"\xc0") -> bytes:
return b"\xff" + marker + struct.pack(">HBHHB", 17, 8, height, width, 3) + b"\x01\x22\x00\x02\x11\x01\x03\x11\x01"
def _jpeg_segment(marker: bytes, payload: bytes) -> bytes:
return b"\xff" + marker + struct.pack(">H", len(payload) + 2) + payload
_EMPTY_JPEG_SEGMENT: Final = _jpeg_segment(b"\xe0", b"")
@pytest.mark.parametrize(
"image",
"marker",
[pytest.param(bytes([marker]), id=hex(marker)) for marker in range(0xC0, 0xD0) if marker not in (0xC4, 0xC8, 0xCC)],
)
def test_image_dimensions_from_bytes_reads_every_start_of_frame_marker(marker: bytes) -> None:
assert image_dimensions_from_bytes(b"\xff\xd8" + _jpeg_sof(800, 600, marker)) == (800, 600)
@pytest.mark.parametrize(
"marker", [pytest.param(b"\xc4", id="dht"), pytest.param(b"\xc8", id="jpg"), pytest.param(b"\xcc", id="dac")]
)
def test_image_dimensions_from_bytes_skips_the_non_frame_markers_in_the_start_of_frame_range(marker: bytes) -> None:
lookalike: Final = _jpeg_segment(marker, struct.pack(">BHH", 8, 1, 1) + b"\x00" * 8)
assert image_dimensions_from_bytes(b"\xff\xd8" + lookalike + _jpeg_sof(800, 600)) == (800, 600)
def test_image_dimensions_from_bytes_skips_each_segment_by_its_own_length() -> None:
segments: Final = (
_jpeg_segment(b"\xe0", b"JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00"),
_jpeg_segment(b"\xe1", b"Exif\x00\x00" + _jpeg_sof(1, 1) + b"\x00" * 7),
_jpeg_segment(b"\xfe", b"c"),
_jpeg_segment(b"\xdb", b"\x00" * 65),
)
assert image_dimensions_from_bytes(b"\xff\xd8" + b"".join(segments) + _jpeg_sof(800, 600)) == (800, 600)
@pytest.mark.parametrize(
("segments_before_frame", "expected"),
[
pytest.param(b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1025 + _jpeg_sof(800, 600), id="too-many-segments"),
pytest.param(b"\xff\xd8" + b"\xff" * 2000 + _jpeg_sof(800, 600)[1:], id="too-many-fill-bytes"),
pytest.param(b"\xff\xd8\xff\xe0\x00\x00\x02" + _jpeg_sof(800, 600), id="segment-length-below-two"),
pytest.param(MAX_JPEG_HEADER_SEGMENTS - 1, (800, 600), id="frame-is-the-last-segment-read"),
pytest.param(MAX_JPEG_HEADER_SEGMENTS, None, id="frame-is-past-the-segment-limit"),
],
)
def test_image_dimensions_from_bytes_gives_up_on_pathological_jpeg_headers(image: bytes) -> None:
assert image_dimensions_from_bytes(image) is None
def test_image_dimensions_from_bytes_reads_at_most_the_segment_limit(
segments_before_frame: int, expected: tuple[int, int] | None
) -> None:
image: Final = b"\xff\xd8" + _EMPTY_JPEG_SEGMENT * segments_before_frame + _jpeg_sof(800, 600)
assert image_dimensions_from_bytes(image) == expected
def test_image_dimensions_from_bytes_still_reads_a_jpeg_with_many_real_segments() -> None:
image: Final = b"\xff\xd8" + b"\xff\xe0\x00\x02" * 1000 + b"\xff" * 64 + _jpeg_sof(800, 600)[1:]
@pytest.mark.parametrize(
("fill_bytes", "expected"),
[
pytest.param(MAX_JPEG_FILL_BYTES - 1, (800, 600), id="marker-is-the-last-byte-read"),
pytest.param(MAX_JPEG_FILL_BYTES, None, id="marker-is-past-the-fill-byte-limit"),
],
)
def test_image_dimensions_from_bytes_reads_at_most_the_fill_byte_limit(
fill_bytes: int, expected: tuple[int, int] | None
) -> None:
image: Final = b"\xff\xd8" + b"\xff" * fill_bytes + _jpeg_sof(800, 600)[1:]
assert image_dimensions_from_bytes(image) == (800, 600)
assert image_dimensions_from_bytes(image) == expected
def test_image_dimensions_from_bytes_gives_up_on_a_segment_length_below_two() -> None:
assert image_dimensions_from_bytes(b"\xff\xd8\xff\xe0\x00\x00\x02" + _jpeg_sof(800, 600)) is None
@pytest.mark.parametrize(
@ -1555,7 +1608,14 @@ def test_get_image_dimensions_still_raises_for_a_truncated_header(header: bytes)
"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"),
pytest.param(
b"\xff\xd8" + _EMPTY_JPEG_SEGMENT * MAX_JPEG_HEADER_SEGMENTS + _jpeg_sof(800, 600),
id="too-many-jpeg-segments",
),
pytest.param(
b"\xff\xd8" + b"\xff" * MAX_JPEG_FILL_BYTES + _jpeg_sof(800, 600)[1:], id="too-many-jpeg-fill-bytes"
),
pytest.param(b"\xff\xd8\xff\xe0\x00\x01\x02" + _jpeg_sof(800, 600), id="jpeg-segment-length-one"),
],
)
def test_get_image_dimensions_falls_back_to_the_default_size_for_a_header_it_cannot_read(image: bytes) -> None:

14
tests/unit/llms/azure_ai/conftest.py Normal file
View file

@ -0,0 +1,14 @@
import logging
from collections.abc import Iterator
import pytest
from litellm._logging import verbose_logger
@pytest.fixture
def litellm_warnings(caplog: pytest.LogCaptureFixture) -> Iterator[pytest.LogCaptureFixture]:
verbose_logger.addHandler(caplog.handler)
with caplog.at_level(logging.WARNING, logger="LiteLLM"):
yield caplog
verbose_logger.removeHandler(caplog.handler)

8
tests/unit/llms/azure_ai/image_edit/test_azure_ai_image_edit_transformation.py
View file

@ -419,12 +419,13 @@ def test_flux2_image_edit_measures_a_stream_reference():
assert response._hidden_params["response_cost"] == pytest.approx(rate * 1024 * 1024 + rate * 2048 * 2048)
# Billable megapixels for a lone reference as Azure's FLUX.2-pro request_meta reported them on 2026-09-25
# Billable megapixels for a lone reference follow Azure's FLUX.2-pro request_meta as reported on 2026-09-25
@pytest.mark.parametrize(
("reference", "billed_megapixels"),
(
pytest.param(_png(640, 640), 1, id="small-reference-rounds-up"),
pytest.param(_png(1024, 1280), 2, id="fractional-reference-rounds-up"),
pytest.param(_png(2048, 1024), 2, id="landscape-reference-is-measured-by-area"),
pytest.param(_jpeg(4032, 3024), 4, id="photo-reference-caps-at-four-megapixels"),
),
)
@ -453,7 +454,9 @@ def test_flux2_image_edit_bills_a_lone_reference_in_whole_megapixels(reference:
pytest.param(_png(0, 640), id="zero-width-header"),
),
)
def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel(reference: bytes):
def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel(
reference: bytes, litellm_warnings: pytest.LogCaptureFixture
):
response: Final = litellm.image_edit(
model="azure_ai/FLUX.2-flex",
image=[reference],
@ -467,6 +470,7 @@ def test_flux2_image_edit_bills_a_lone_unmeasurable_reference_as_one_megapixel(r
assert response._hidden_params["reference_image_pixels"] == (UNMEASURED_REFERENCE_IMAGE_PIXELS,)
assert response._hidden_params["response_cost"] == pytest.approx(rate * 1024 * 1024 + rate * 1024 * 1024)
assert "billing it as one megapixel" in litellm_warnings.text
def test_flux2_image_edit_still_bills_every_reference_when_one_header_reports_zero_pixels():

48
tests/unit/llms/azure_ai/image_generation/test_azure_ai_flux2_image_generation.py
View file

@ -405,6 +405,32 @@ def test_flux2_pro_bills_one_megapixel_for_a_size_it_cannot_measure(size: str):
assert cost == pytest.approx(first)
@pytest.mark.parametrize(
("optional_params", "megapixels"),
(
pytest.param({"width": 1, "height": 1}, 1, id="smallest-mapped-dimensions-win"),
pytest.param({"width": 0, "height": 1024}, 2, id="zero-width-falls-back-to-size"),
pytest.param({"width": 2048, "height": 0}, 2, id="zero-height-falls-back-to-size"),
pytest.param({"width": True, "height": 1024}, 2, id="bool-width-falls-back-to-size"),
),
)
def test_flux2_pro_bills_mapped_dimensions_only_when_both_are_positive_integers(
optional_params: dict[str, int | bool], megapixels: int
):
first, additional = _pro_megapixel_rates()
cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
model="flux.2-pro",
completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]),
custom_llm_provider="azure_ai",
size="2048x1024",
optional_params=optional_params,
call_type="image_generation",
)
assert cost == pytest.approx(first + additional * (megapixels - 1))
def test_flux2_pro_bills_the_requested_size_when_the_returned_image_is_not_valid_base64():
first, additional = _pro_megapixel_rates()
@ -435,14 +461,15 @@ def test_flux2_pro_bills_the_requested_image_count_when_the_response_lists_none(
assert cost == pytest.approx(first * billed_images)
def test_flux2_pro_reads_an_uppercase_size():
@pytest.mark.parametrize("size", ("2048X1024", "2048-x-1024"))
def test_flux2_pro_reads_each_size_spelling(size: str):
first, additional = _pro_megapixel_rates()
cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
model="flux.2-pro",
completion_response=ImageResponse(data=[ImageObject(b64_json="aW1n")]),
custom_llm_provider="azure_ai",
size="2048X1024",
size=size,
call_type="image_generation",
)
@ -511,6 +538,7 @@ def test_flux2_flex_edit_bills_reference_pixels_on_top_of_generated_pixels():
assert _flex_edit_cost(_edit_response((3 * 1024 * 1024,))) - generated_only == pytest.approx(
_flex_pixel_rate() * 3 * 1024 * 1024
)
assert _flex_edit_cost(_edit_response((1,))) - generated_only == pytest.approx(_flex_pixel_rate() * 1024 * 1024)
def test_flux2_flex_edit_reference_cost_does_not_scale_with_image_count():
@ -535,6 +563,15 @@ def test_flux2_flex_edit_bills_only_positive_integer_reference_counts(reference_
assert _flex_edit_cost(_edit_response(reference_image_pixels)) == generated_only
@pytest.mark.parametrize("reference_image_pixels", ((0,), ("1048576",), 1048576))
def test_flux2_flex_edit_warns_when_it_ignores_malformed_reference_counts(
reference_image_pixels: object, litellm_warnings: pytest.LogCaptureFixture
):
_flex_edit_cost(_edit_response(reference_image_pixels))
assert "Ignoring malformed FLUX.2 reference pixel counts" in litellm_warnings.text
def test_flux2_flex_edit_ignores_reference_pixels_when_provider_reports_token_usage():
response: Final = ImageResponse(
data=[ImageObject(b64_json="aW1n")],
@ -634,19 +671,20 @@ def test_flux2_flex_cost_prefers_deployment_input_cost_per_pixel() -> None:
assert cost == pytest.approx(2e-07 * 2048 * 1024 * 2)
def test_unlisted_azure_ai_model_bills_deployment_input_cost_per_pixel() -> None:
@pytest.mark.parametrize(("width", "height"), ((1024, 1024), (1536, 1024)), ids=("whole-megapixel", "fractional"))
def test_unlisted_azure_ai_model_bills_deployment_input_cost_per_pixel(width: int, height: int) -> None:
response: Final = ImageResponse(data=[ImageObject(b64_json="aW1n"), ImageObject(b64_json="aW1n")])
cost: Final = CostCalculatorUtils.route_image_generation_cost_calculator(
model="unlisted-flux-deployment",
completion_response=response,
custom_llm_provider="azure_ai",
size="1024x1024",
size=f"{width}x{height}",
call_type="image_generation",
model_info={"input_cost_per_pixel": 1e-07},
)
assert cost == pytest.approx(1e-07 * 1024 * 1024 * 2)
assert cost == pytest.approx(1e-07 * width * height * 2)
def test_unlisted_azure_ai_deployment_without_a_generated_image_price_bills_nothing() -> None: