From 2f796530d06fef5e5a7aa0f007b772de274b6b8d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:50:51 -0700 Subject: [PATCH] test(health): parse the probe PNG without mutation --- .../test_health_check_helpers.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py index a5f82cf04a6..ee2a31beff7 100644 --- a/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_health_check_helpers.py @@ -17,18 +17,19 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS +def _png_chunks(png: bytes, offset: int = 8) -> tuple[tuple[bytes, bytes], ...]: + if offset >= len(png): + return () + (length,) = struct.unpack(">I", png[offset : offset + 4]) + chunk = (png[offset + 4 : offset + 8], png[offset + 8 : offset + 8 + length]) + return (chunk, *_png_chunks(png, offset + 12 + length)) + + def _distinct_rgb_colors(png: bytes) -> set[bytes]: width = int.from_bytes(png[16:20], "big") - idat_parts = [] - offset = 8 - while offset < len(png): - (length,) = struct.unpack(">I", png[offset : offset + 4]) - if png[offset + 4 : offset + 8] == b"IDAT": - idat_parts.append(png[offset + 8 : offset + 8 + length]) - offset += 12 + length - raw = zlib.decompress(b"".join(idat_parts)) + raw = zlib.decompress(b"".join(data for tag, data in _png_chunks(png) if tag == b"IDAT")) row_size = 1 + width * 3 - rows = [raw[i : i + row_size] for i in range(0, len(raw), row_size)] + rows = tuple(raw[i : i + row_size] for i in range(0, len(raw), row_size)) assert all(row[0] == 0 for row in rows) return {bytes(row[i : i + 3]) for row in rows for i in range(1, row_size, 3)}