diff --git a/doc/api.rst b/doc/api.rst index 145f729cf3..d64c75e566 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -307,6 +307,7 @@ API Reference FuzzerShortenConverter FuzzerSimilarConverter HumanInTheLoopConverter + ImageCompressionConverter InsertPunctuationConverter LeetspeakConverter LLMGenericTextConverter diff --git a/pyrit/prompt_converter/__init__.py b/pyrit/prompt_converter/__init__.py index 176218c3c8..87c4e29fda 100644 --- a/pyrit/prompt_converter/__init__.py +++ b/pyrit/prompt_converter/__init__.py @@ -36,6 +36,7 @@ FuzzerSimilarConverter, ) from pyrit.prompt_converter.human_in_the_loop_converter import HumanInTheLoopConverter +from pyrit.prompt_converter.image_compression_converter import ImageCompressionConverter from pyrit.prompt_converter.insert_punctuation_converter import InsertPunctuationConverter from pyrit.prompt_converter.leetspeak_converter import LeetspeakConverter from pyrit.prompt_converter.llm_generic_text_converter import LLMGenericTextConverter @@ -100,6 +101,7 @@ "FuzzerShortenConverter", "FuzzerSimilarConverter", "HumanInTheLoopConverter", + "ImageCompressionConverter", "InsertPunctuationConverter", "TemplateSegmentConverter", "TextJailbreakConverter", diff --git a/pyrit/prompt_converter/image_compression_converter.py b/pyrit/prompt_converter/image_compression_converter.py new file mode 100644 index 0000000000..358f4e9731 --- /dev/null +++ b/pyrit/prompt_converter/image_compression_converter.py @@ -0,0 +1,270 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +import base64 +import logging +from io import BytesIO +from typing import Literal, Optional +from urllib.parse import urlparse + +import aiohttp +from PIL import Image + +from pyrit.models import PromptDataType, data_serializer_factory +from pyrit.prompt_converter import ConverterResult, PromptConverter + +logger = logging.getLogger(__name__) + + +class ImageCompressionConverter(PromptConverter): + """ + Compresses images to reduce file size while preserving visual quality. + + This converter supports multiple compression strategies across JPEG, PNG, and WEBP formats, + each with format-specific optimization settings. It can maintain the original image format + or convert between formats as needed. + + When converting images with transparency (alpha channel) to JPEG format, the converter + automatically composites the transparent areas onto a solid background color. + + Supported input types: + File paths to any image that PIL can open (or URLs pointing to such images): + https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#fully-supported-formats + + Supported output formats and their compression options: + - JPEG: ``quality``, ``optimize``, ``progressive`` + - PNG: ``optimize``, ``compress_level`` + - WEBP: ``quality``, ``lossless``, ``method`` + + References: + https://pillow.readthedocs.io/en/stable/handbook/concepts.html + https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#jpeg-saving + https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#png-saving + https://pillow.readthedocs.io/en/stable/handbook/image-file-formats.html#webp-saving + """ + + def __init__( + self, + *, + output_format: Optional[Literal["JPEG", "PNG", "WEBP"]] = None, + quality: Optional[int] = None, + optimize: Optional[bool] = None, + progressive: Optional[bool] = None, + compress_level: Optional[int] = None, + lossless: Optional[bool] = None, + method: Optional[int] = None, + background_color: tuple[int, int, int] = (0, 0, 0), + min_compression_threshold: int = 1024, + fallback_to_original: bool = True, + ): + """ + Initializes the converter with specified compression settings. + + Args: + output_format (str, optional): Output image format. If None, keeps original format (if supported). + quality (int, optional): General quality setting for JPEG and WEBP formats (0-100).\n + For JPEG format, it represents the image quality, on a scale from 0 (worst) to 95 (best).\n + For WEBP format, the value ranges from 0 to 100; for lossy compression: 0-smallest file size and + 100-largest; for ``lossless``: 0-fastest/less efficient, and 100 gives the best compression. + optimize (bool, optional): Whether to optimize the image during compression. \n + For JPEG: makes the encoder perform an extra pass over the image to select optimal settings.\n + For PNG: instructs the PNG writer to make the output file as small as possible. + progressive (bool, optional): Whether to save JPEG images as progressive. + compress_level (int, optional): ZLIB compression level (0-9): 1=fastest, 9=best, 0=none. + Ignored if ``optimize`` is True (then it is forced to 9). + lossless (bool, optional): Whether to use lossless compression for WEBP format. + method (int, optional): Quality/speed trade-off for WEBP format (0=fast, 6=slower-better). + background_color (tuple[int, int, int]): RGB color tuple for background when converting + transparent images to JPEG. Defaults to black. + min_compression_threshold (int): Minimum file size threshold for compression. Defaults to 1024 bytes. + fallback_to_original (bool): Fallback to original if compression increases file size. Defaults to True. + + Raises: + ValueError: If unsupported output format is specified, or if some of the parameters are out of range. + """ + if quality is not None and (quality < 0 or quality > 100): + raise ValueError("Quality must be between 0 and 100") + self._quality = quality + + if output_format and output_format not in ("JPEG", "PNG", "WEBP"): + raise ValueError("Output format must be one of 'JPEG', 'PNG', or 'WEBP'") + self._output_format = output_format + + if compress_level is not None and (compress_level < 0 or compress_level > 9): + raise ValueError("Compress level must be between 0 and 9") + self._compress_level = compress_level + + if method is not None and (method < 0 or method > 6): + raise ValueError("Method must be between 0 and 6 for WEBP format") + self._method = method + + if min_compression_threshold < 0: + raise ValueError("Minimum compression threshold must be a non-negative integer") + self._min_compression_threshold = min_compression_threshold + + if ( + not isinstance(background_color, tuple) + or len(background_color) != 3 + or not all(isinstance(c, int) and 0 <= c <= 255 for c in background_color) + ): + raise ValueError("Background color must be a tuple of three integers between 0 and 255") + + self._optimize = optimize + self._progressive = progressive + self._lossless = lossless + self._background_color = background_color + self._fallback_to_original = fallback_to_original + + if output_format == "JPEG" and quality is not None and quality > 95: + logger.warning( + "Using quality > 95 for JPEG may result in larger files. Consider using a lower quality setting." + ) + + def _should_compress(self, original_size: int) -> bool: + """Determines if image should be compressed.""" + if original_size < self._min_compression_threshold: + return False # skip compression for small images + return True + + def _compress_image(self, image: Image.Image, original_format: str, original_size: int) -> tuple[BytesIO, str]: + """Compresses the image with the specified settings. Returns the compressed image bytes and output format.""" + original_format = original_format.upper() + output_format = self._output_format or ( + original_format if original_format in ("JPEG", "PNG", "WEBP") else "JPEG" + ) + + logger.info( + f"Compressing image: original format={original_format}, " + f"output format={output_format}, original size={original_size} bytes" + ) + + # Handle images with transparency when converting to JPEG + if output_format == "JPEG": + if image.has_transparency_data: + image = image.convert("RGBA") + background = Image.new("RGB", image.size, self._background_color) + background.paste(image, mask=image.split()[-1]) + image = background + else: + image = image.convert("RGB") + + save_kwargs: dict = {} + + # Format-specific options for currently supported output types + if output_format == "JPEG": + save_kwargs = {} + if self._quality is not None: + save_kwargs["quality"] = self._quality + if self._optimize: + save_kwargs["optimize"] = True + if self._progressive: + save_kwargs["progressive"] = True + elif output_format == "PNG": + save_kwargs = {} + if self._optimize: + save_kwargs["optimize"] = True + if self._compress_level is not None: + save_kwargs["compress_level"] = self._compress_level + elif output_format == "WEBP": + save_kwargs = {} + if self._quality is not None: + save_kwargs["quality"] = self._quality + if self._lossless: + save_kwargs["lossless"] = True + if self._method is not None: + save_kwargs["method"] = self._method + + compressed_bytes = BytesIO() # in-memory buffer + image.save(compressed_bytes, output_format, **save_kwargs) + return compressed_bytes, output_format + + async def _handle_original_image_fallback( + self, prompt: str, input_type: PromptDataType, img_serializer, original_img_bytes: bytes, original_format: str + ) -> ConverterResult: + """Handles fallback to original image for both URL and file path inputs.""" + if input_type == "url": + # We need to save the downloaded content locally and return the local path + img_serializer.file_extension = original_format.lower() + await img_serializer.save_data(original_img_bytes) + return ConverterResult(output_text=str(img_serializer.value), output_type="image_path") + return ConverterResult(output_text=prompt, output_type="image_path") + + async def _read_image_from_url(self, url: str) -> bytes: + """Downloads data from URL and returns the content as bytes.""" + try: + async with aiohttp.ClientSession() as session: + async with session.get(url) as response: + response.raise_for_status() + return await response.read() + except aiohttp.ClientError as e: + raise RuntimeError(f"Failed to download content from URL {url}: {str(e)}") + + async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: + """ + Converts the given prompt (image) by compressing it. + + Args: + prompt (str): The image file path or URL pointing to the image to be compressed. + input_type (PromptDataType): The type of input data. + + Returns: + ConverterResult: The result containing path to the compressed image. + + Raises: + ValueError: If the input type is not supported. + """ + if input_type not in ("image_path", "url"): + raise ValueError(f"Input type '{input_type}' not supported") + if input_type == "url" and urlparse(prompt).scheme not in ("http", "https"): + raise ValueError(f"Invalid URL: {prompt}. Must start with 'http://' or 'https://'") + + img_serializer = data_serializer_factory(category="prompt-memory-entries", value=prompt, data_type="image_path") + + # Read the image data into memory as bytes for processing + original_img_bytes = ( + await self._read_image_from_url(prompt) if input_type == "url" else await img_serializer.read_data() + ) + original_img = Image.open(BytesIO(original_img_bytes)) + + original_format = original_img.format or "JPEG" # since PIL may not always provide a format + original_size = len(original_img_bytes) + + # This is to avoid unnecessary processing and potential quality loss for images that are already small + if not self._should_compress(original_size): + logger.warning(f"Image too small ({original_size} bytes), skipping compression") + return await self._handle_original_image_fallback( + prompt, input_type, img_serializer, original_img_bytes, original_format + ) + + # Compress the image and get back a BytesIO buffer containing the compressed data + # along with the actual output format used (which may differ from input format) + compressed_bytes, output_format = self._compress_image(original_img, original_format, original_size) + compressed_bytes_value = compressed_bytes.getvalue() + compressed_size = len(compressed_bytes_value) + + # Sometimes compression can actually increase file size so we check if we should fallback to the original + if self._fallback_to_original and compressed_size >= original_size: + logger.warning(f"Compression increased file size ({original_size} → {compressed_size}), using original") + return await self._handle_original_image_fallback( + prompt, input_type, img_serializer, original_img_bytes, original_format + ) + + # This ensures the saved file has the correct extension for its actual format + # Only currently supported output formats are taken into account + format_extensions = {"JPEG": "jpeg", "PNG": "png", "WEBP": "webp"} + img_serializer.file_extension = format_extensions.get(output_format, "jpeg") + + # Convert compressed bytes to base64 for storage via the serializer + image_str = base64.b64encode(compressed_bytes_value) + await img_serializer.save_b64_image(data=image_str.decode()) + + compression_ratio = (1 - compressed_size / original_size) * 100 if original_size > 0 else 0 + logger.info(f"Image compressed: {original_size} → {compressed_size} ({compression_ratio:.1f}% reduction)") + + return ConverterResult(output_text=str(img_serializer.value), output_type="image_path") + + def input_supported(self, input_type: PromptDataType) -> bool: + return input_type in ("image_path", "url") + + def output_supported(self, output_type: PromptDataType) -> bool: + return output_type == "image_path" diff --git a/tests/unit/converter/test_image_compression_converter.py b/tests/unit/converter/test_image_compression_converter.py new file mode 100644 index 0000000000..ffe3afc3fb --- /dev/null +++ b/tests/unit/converter/test_image_compression_converter.py @@ -0,0 +1,390 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from io import BytesIO +from unittest.mock import AsyncMock, patch + +import pytest +from PIL import Image + +from pyrit.prompt_converter import ImageCompressionConverter + + +@pytest.fixture +def sample_image_bytes(): + """Sample RGB image for testing with configurable format and size.""" + + def _create_image(format="PNG", size=(200, 200)): + img = Image.new("RGB", size, color=(125, 125, 125)) + img_bytes = BytesIO() + img.save(img_bytes, format=format) + return img_bytes.getvalue() + + return _create_image + + +@pytest.fixture +def sample_transparent_image_bytes(): + """Sample RGBA image with transparency for testing with configurable format.""" + + def _create_image(format="PNG"): + img = Image.new("RGBA", (200, 200), color=(125, 125, 125, 128)) + img_bytes = BytesIO() + img.save(img_bytes, format=format) + return img_bytes.getvalue() + + return _create_image + + +def test_image_compression_converter_initialization_output_format_validation(): + """Test validation of output_format parameter.""" + for unsupported_format in ["GIF", "BMP", "TIFF", "ICO", "WEBM", "SVG", "jpg", "png"]: + with pytest.raises(ValueError, match="Output format must be one of 'JPEG', 'PNG', or 'WEBP'"): + ImageCompressionConverter(output_format=unsupported_format) # type: ignore + + supported_formats = ["JPEG", "PNG", "WEBP"] + for supported_format in supported_formats: + converter = ImageCompressionConverter(output_format=supported_format) # type: ignore + assert converter._output_format == supported_format + + converter = ImageCompressionConverter(output_format=None) + assert converter._output_format is None + + +def test_image_compression_converter_initialization_quality_validation(): + """Test validation of quality parameter.""" + for invalid_quality in [-1, -10, 101, 150, 999]: + with pytest.raises(ValueError, match="Quality must be between 0 and 100"): + ImageCompressionConverter(quality=invalid_quality) + + for valid_quality in [0, 1, 50, 95, 100]: + converter = ImageCompressionConverter(quality=valid_quality) + assert converter._quality == valid_quality + + converter = ImageCompressionConverter(quality=None) + assert converter._quality is None + + +def test_image_compression_converter_initialization_compress_level_validation(): + """Test validation of compress_level parameter.""" + for invalid_level in [-1, -5, 10, 15, 100]: + with pytest.raises(ValueError, match="Compress level must be between 0 and 9"): + ImageCompressionConverter(compress_level=invalid_level) + + for valid_level in [0, 1, 5, 9]: + converter = ImageCompressionConverter(compress_level=valid_level) + assert converter._compress_level == valid_level + + converter = ImageCompressionConverter(compress_level=None) + assert converter._compress_level is None + + +def test_image_compression_converter_initialization_method_validation(): + """Test validation of method parameter for WEBP format.""" + for invalid_method in [-1, -5, 7, 10, 100]: + with pytest.raises(ValueError, match="Method must be between 0 and 6 for WEBP format"): + ImageCompressionConverter(method=invalid_method) + + for valid_method in [0, 1, 3, 6]: + converter = ImageCompressionConverter(method=valid_method) + assert converter._method == valid_method + + converter = ImageCompressionConverter(method=None) + assert converter._method is None + + +def test_image_compression_converter_initialization_min_compression_threshold_validation(): + """Test validation of min_compression_threshold parameter.""" + for invalid_threshold in [-1, -10, -100, -1024]: + with pytest.raises(ValueError, match="Minimum compression threshold must be a non-negative integer"): + ImageCompressionConverter(min_compression_threshold=invalid_threshold) + + for valid_threshold in [0, 1, 512, 1024, 2048]: + converter = ImageCompressionConverter(min_compression_threshold=valid_threshold) + assert converter._min_compression_threshold == valid_threshold + + +def test_image_compression_converter_initialization_background_color_validation(): + """Test validation of background_color parameter.""" + invalid_colors = [ + "black", + [0, 0, 0], + (0, 0), + (0, 0, 256), + (-0.5, 0, 0), + (None, 0, 0), + ] + + for invalid_color in invalid_colors: + with pytest.raises(ValueError, match="Background color must be a tuple of three integers between 0 and 255"): + ImageCompressionConverter(background_color=invalid_color) # type: ignore + + valid_colors = [ + (0, 0, 0), + (255, 255, 255), + (255, 0, 0), + (0, 255, 0), + (0, 0, 255), + (128, 128, 128), + ] + + for valid_color in valid_colors: + converter = ImageCompressionConverter(background_color=valid_color) + assert converter._background_color == valid_color + + +def test_image_compression_converter_initialization_valid_combinations(): + """Test that valid parameter combinations work correctly.""" + converter = ImageCompressionConverter( + output_format="JPEG", + quality=85, + optimize=True, + progressive=True, + compress_level=6, + lossless=False, + method=4, + background_color=(255, 255, 255), + min_compression_threshold=2048, + fallback_to_original=False, + ) + + assert converter._output_format == "JPEG" + assert converter._quality == 85 + assert converter._optimize is True + assert converter._progressive is True + assert converter._compress_level == 6 + assert converter._lossless is False + assert converter._method == 4 + assert converter._background_color == (255, 255, 255) + assert converter._min_compression_threshold == 2048 + assert converter._fallback_to_original is False + + +def test_image_compression_converter_quality_warning(): + """Test that high quality values for JPEG trigger a warning.""" + with patch("pyrit.prompt_converter.image_compression_converter.logger") as mock_logger: + ImageCompressionConverter(output_format="JPEG", quality=98) + mock_logger.warning.assert_called_once() + + +@pytest.mark.parametrize( + "input_format, output_format, expected_output_format", + [ + # Format preservation + ("JPEG", None, "JPEG"), + ("WEBP", None, "WEBP"), + # Cross-format conversion + ("PNG", "JPEG", "JPEG"), + ("WEBP", "PNG", "PNG"), + ], +) +@pytest.mark.asyncio +async def test_image_compression_converter_format_preservation_and_conversion( + duckdb_instance, sample_image_bytes, input_format, output_format, expected_output_format +): + """Test format preservation and conversion between formats.""" + converter = ImageCompressionConverter( + output_format=output_format, min_compression_threshold=100, fallback_to_original=False + ) + # Create larger image to ensure compression happens + image_bytes = sample_image_bytes(format=input_format, size=(2048, 2048)) + + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_serializer.read_data.return_value = image_bytes + mock_serializer.save_b64_image = AsyncMock() + # Set the value to match input format initially + mock_serializer.value = f"test_image.{input_format.lower()}" + # Mock the file_extension property to be settable + mock_serializer.file_extension = input_format.lower() + mock_factory.return_value = mock_serializer + + await converter.convert_async(prompt=f"test_image.{input_format.lower()}", input_type="image_path") + + # Verify the save method was called + mock_serializer.save_b64_image.assert_called_once() + mock_serializer.read_data.assert_called_once() + + # Check that file extension was updated correctly + expected_extension = expected_output_format.lower() + assert mock_serializer.file_extension == expected_extension + + +@pytest.mark.parametrize( + "input_format, output_format, background_color, expected_output_format", + [ + ("PNG", "JPEG", (255, 255, 255), "JPEG"), + ("WEBP", "JPEG", (0, 0, 0), "JPEG"), + ("TIFF", "JPEG", (255, 0, 0), "JPEG"), + ], +) +@pytest.mark.asyncio +async def test_image_compression_converter_transparency_handling( + duckdb_instance, + sample_transparent_image_bytes, + input_format, + output_format, + background_color, + expected_output_format, +): + """Test transparency handling across formats with different background colors.""" + converter = ImageCompressionConverter(output_format=output_format, background_color=background_color) + image_bytes = sample_transparent_image_bytes(format=input_format) + image = Image.open(BytesIO(image_bytes)) + + assert image.has_transparency_data # before compression, the image should have transparency + + res_compressed_io, res_output_format = converter._compress_image(image, input_format, len(image_bytes)) + assert res_compressed_io + assert res_output_format == expected_output_format + + output_image = Image.open(res_compressed_io) + assert output_image.has_transparency_data is False # after compression, the image should not have transparency + + +@pytest.mark.asyncio +async def test_image_compression_converter_skip(duckdb_instance, sample_image_bytes): + """Test cases for skipping compression and fallback to original image.""" + + # 1: Skip compression for small images + converter = ImageCompressionConverter(min_compression_threshold=4567) + small_image_bytes = sample_image_bytes(format="PNG", size=(50, 50)) + assert len(small_image_bytes) < 4567 + + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_serializer.read_data.return_value = small_image_bytes + mock_factory.return_value = mock_serializer + + result = await converter.convert_async(prompt="small_image.png", input_type="image_path") + assert result.output_text == "small_image.png" + + # Verify that compression was skipped - save_b64_image should not be called + mock_serializer.save_b64_image.assert_not_called() + + # 2: Fallback to original when compression increases file size + converter_fallback = ImageCompressionConverter(fallback_to_original=True, quality=100) + large_image_bytes = sample_image_bytes(format="PNG", size=(500, 500)) + + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_serializer.read_data.return_value = large_image_bytes + mock_factory.return_value = mock_serializer + + # Mock _compress_image to return larger size + with patch.object(converter_fallback, "_compress_image") as mock_compress: + larger_bytes = BytesIO(b"x" * (len(large_image_bytes) + 1000)) + mock_compress.return_value = (larger_bytes, "PNG") + + result = await converter_fallback.convert_async(prompt="test.png", input_type="image_path") + assert result.output_text == "test.png" + + mock_serializer.save_b64_image.assert_not_called() + mock_compress.assert_called_once() # compression was attempted but resulted in larger size + + +@pytest.mark.asyncio +async def test_image_compression_converter_actually_compresses(duckdb_instance, sample_image_bytes): + """Test that compression actually reduces file size for appropriate images.""" + converter = ImageCompressionConverter(compress_level=9) + + # Create a large image that should benefit from compression + large_image_bytes = sample_image_bytes(format="PNG", size=(1024, 1024)) + original_size = len(large_image_bytes) + image = Image.open(BytesIO(large_image_bytes)) + + compressed_io, output_format = converter._compress_image(image, "PNG", original_size) + compressed_size = len(compressed_io.getvalue()) + + assert compressed_size < original_size + + +@pytest.mark.asyncio +async def test_image_compression_converter_url_format_conversion(duckdb_instance, sample_image_bytes): + """Test successful compression of image from URL.""" + converter = ImageCompressionConverter(output_format="WEBP", min_compression_threshold=100) + test_url = "https://example.com/test_image.jpeg" + large_image_bytes = sample_image_bytes(format="JPEG", size=(2048, 2048)) + + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_serializer.file_extension = "jpeg" + mock_serializer.value = "converted_image.webp" + mock_serializer.save_b64_image = AsyncMock() + mock_factory.return_value = mock_serializer + + with patch.object(converter, "_read_image_from_url") as mock_read_url: + mock_read_url.return_value = large_image_bytes + + result = await converter.convert_async(prompt=test_url, input_type="url") + + assert result.output_text == "converted_image.webp" + assert result.output_type == "image_path" + # Verify file extension was updated to match WEBP output format + assert mock_serializer.file_extension == "webp" + mock_serializer.save_b64_image.assert_called_once() + + +@pytest.mark.asyncio +async def test_image_compression_converter_url_input_fallback_scenarios(duckdb_instance, sample_image_bytes): + """Test URL input fallback scenarios (small image size).""" + converter = ImageCompressionConverter(min_compression_threshold=5000, fallback_to_original=True) + test_url = "https://example.com/small_image.png" + small_image_bytes = sample_image_bytes(format="PNG", size=(100, 100)) + + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_serializer.file_extension = "png" + mock_serializer.value = "fallback_image.png" + mock_serializer.save_data = AsyncMock() + mock_factory.return_value = mock_serializer + + with patch.object(converter, "_read_image_from_url") as mock_read_url: + mock_read_url.return_value = small_image_bytes + + result = await converter.convert_async(prompt=test_url, input_type="url") + + assert result.output_text == "fallback_image.png" + assert result.output_type == "image_path" + mock_read_url.assert_called_once_with(test_url) + mock_serializer.save_data.assert_called_once_with(small_image_bytes) + mock_serializer.save_b64_image.assert_not_called() + + +@pytest.mark.asyncio +async def test_image_compression_converter_invalid_url(): + """Test handling of invalid URLs.""" + converter = ImageCompressionConverter() + invalid_urls = ["ftp://example.com/image.png", "file:///local/path/image.png", "not-url", "example.com/image.png"] + for invalid_url in invalid_urls: + with pytest.raises(ValueError, match="Invalid URL"): + await converter.convert_async(prompt=invalid_url, input_type="url") + + +@pytest.mark.asyncio +async def test_image_compression_converter_corrupted_image_bytes(): + """Test handling of corrupted image bytes.""" + converter = ImageCompressionConverter() + corrupted_bytes = b"notanimagefile" + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_serializer.read_data.return_value = corrupted_bytes + mock_factory.return_value = mock_serializer + with pytest.raises(Exception): + await converter.convert_async(prompt="corrupted.png", input_type="image_path") + + +@pytest.mark.asyncio +async def test_image_compression_converter_output_format_fallback(sample_image_bytes): + """Test fallback to JPEG when original format is unsupported (and no output_format specified).""" + img = Image.new("RGB", (100, 100), color=(123, 123, 123)) + img_bytes = BytesIO() + img.save(img_bytes, format="TIFF") + img_bytes = img_bytes.getvalue() + converter = ImageCompressionConverter(output_format=None) + with patch("pyrit.prompt_converter.image_compression_converter.data_serializer_factory") as mock_factory: + mock_serializer = AsyncMock() + mock_factory.return_value = mock_serializer + mock_serializer.read_data.return_value = img_bytes + await converter.convert_async(prompt="test.tiff", input_type="image_path") + assert mock_serializer.file_extension == "jpeg"