From b1b70ca5a98152934d36eaea7bb1470f825f93f2 Mon Sep 17 00:00:00 2001 From: nightcityblade Date: Wed, 19 Aug 2026 11:29:27 +0800 Subject: [PATCH 1/2] MAINT Consolidate image URL download handling --- .../base_image_to_image_converter.py | 16 ++++++++----- .../converter/image_compression_converter.py | 9 ++----- .../test_image_compression_converter.py | 24 +++++++++++++++++++ 3 files changed, 36 insertions(+), 13 deletions(-) diff --git a/pyrit/converter/base_image_to_image_converter.py b/pyrit/converter/base_image_to_image_converter.py index 1a2a172c5f..24c4ca24e1 100644 --- a/pyrit/converter/base_image_to_image_converter.py +++ b/pyrit/converter/base_image_to_image_converter.py @@ -18,6 +18,15 @@ logger = logging.getLogger(__name__) +async def _download_image_from_url_async(url: str) -> bytes: + try: + async with aiohttp.ClientSession() as session, 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)}") from e + + class BaseImageToImageConverter(Converter, ABC): """ Abstract base class for image converters that apply a transformation to an image. @@ -139,12 +148,7 @@ async def _read_image_from_url_async(self, url: str) -> bytes: Raises: RuntimeError: If there is an error during the download process. """ - try: - async with aiohttp.ClientSession() as session, 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)}") from e + return await _download_image_from_url_async(url) async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: """ diff --git a/pyrit/converter/image_compression_converter.py b/pyrit/converter/image_compression_converter.py index 43059a87d0..35b0c575d4 100644 --- a/pyrit/converter/image_compression_converter.py +++ b/pyrit/converter/image_compression_converter.py @@ -7,9 +7,9 @@ from typing import Any, Literal from urllib.parse import urlparse -import aiohttp from PIL import Image +from pyrit.converter.base_image_to_image_converter import _download_image_from_url_async from pyrit.converter.converter import Converter, ConverterResult from pyrit.memory import data_serializer_factory from pyrit.models import ComponentIdentifier, PromptDataType @@ -259,12 +259,7 @@ async def _read_image_from_url_async(self, url: str) -> bytes: Raises: RuntimeError: If there is an error during the download process. """ - try: - async with aiohttp.ClientSession() as session, 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)}") from e + return await _download_image_from_url_async(url) async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult: """ diff --git a/tests/unit/converter/test_image_compression_converter.py b/tests/unit/converter/test_image_compression_converter.py index 76165e246b..1bb75599da 100644 --- a/tests/unit/converter/test_image_compression_converter.py +++ b/tests/unit/converter/test_image_compression_converter.py @@ -1,13 +1,37 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT license. +import asyncio from io import BytesIO from unittest.mock import AsyncMock, patch +import aiohttp import pytest from PIL import Image from pyrit.converter import ImageCompressionConverter +from pyrit.converter.base_image_to_image_converter import _download_image_from_url_async + + +async def test_download_image_from_url_async_preserves_response_semantics(): + response = AsyncMock() + with patch("pyrit.converter.base_image_to_image_converter.aiohttp.ClientSession") as client: + response.raise_for_status = client.raise_for_status + session = client.session + client.return_value.__aenter__.return_value = session + session.get.return_value.__aenter__.return_value = response + + response.read.return_value = b"image" + assert await _download_image_from_url_async("https://example.com/image") == b"image" + + response.raise_for_status.side_effect = aiohttp.ClientResponseError(client, (), status=404) + with pytest.raises(RuntimeError, match="Failed to download content from URL"): + await _download_image_from_url_async("https://example.com/image") + + response.raise_for_status.side_effect = None + response.read.side_effect = asyncio.CancelledError() + with pytest.raises(asyncio.CancelledError): + await _download_image_from_url_async("https://example.com/image") @pytest.fixture From 7f3de9ea513cc1b3aca950c32b4163a5af150e3b Mon Sep 17 00:00:00 2001 From: nightcityblade Date: Wed, 19 Aug 2026 23:07:37 +0800 Subject: [PATCH 2/2] TEST Cover shared image download delegation --- .../test_image_compression_converter.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/unit/converter/test_image_compression_converter.py b/tests/unit/converter/test_image_compression_converter.py index 1bb75599da..545263496a 100644 --- a/tests/unit/converter/test_image_compression_converter.py +++ b/tests/unit/converter/test_image_compression_converter.py @@ -9,7 +9,7 @@ import pytest from PIL import Image -from pyrit.converter import ImageCompressionConverter +from pyrit.converter import ImageCompressionConverter, ImageRotationConverter from pyrit.converter.base_image_to_image_converter import _download_image_from_url_async @@ -34,6 +34,24 @@ async def test_download_image_from_url_async_preserves_response_semantics(): await _download_image_from_url_async("https://example.com/image") +async def test_base_image_converter_delegates_url_download(): + with patch( + "pyrit.converter.base_image_to_image_converter._download_image_from_url_async", + new=AsyncMock(return_value=b"image"), + ) as download: + assert await ImageRotationConverter()._read_image_from_url_async("https://example.com/image") == b"image" + download.assert_awaited_once_with("https://example.com/image") + + +async def test_image_compression_converter_delegates_url_download(): + with patch( + "pyrit.converter.image_compression_converter._download_image_from_url_async", + new=AsyncMock(return_value=b"image"), + ) as download: + assert await ImageCompressionConverter()._read_image_from_url_async("https://example.com/image") == b"image" + download.assert_awaited_once_with("https://example.com/image") + + @pytest.fixture def sample_image_bytes(): """Sample RGB image for testing with configurable format and size."""