MAINT Consolidate image URL download handling

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:
         """

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:
         """

tests/unit/converter/test_image_compression_converter.py

@@ -1,13 +1,55 @@
 # 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 import ImageCompressionConverter, ImageRotationConverter
+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")
+
+
+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