Skip to content

Commit 72c0906

Browse files
committed
MAINT: consolidate duplicate image URL download handling
1 parent 62aa4da commit 72c0906

4 files changed

Lines changed: 217 additions & 14 deletions

File tree

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT license.
3+
4+
import aiohttp
5+
6+
7+
async def _download_image_from_url_async(url: str) -> bytes:
8+
"""
9+
Download image data from a URL and return it as bytes.
10+
11+
Args:
12+
url (str): The URL to download the image from.
13+
14+
Returns:
15+
bytes: The content of the image as bytes.
16+
17+
Raises:
18+
RuntimeError: If an aiohttp client error occurs during the download process.
19+
"""
20+
try:
21+
async with aiohttp.ClientSession() as session, session.get(url) as response:
22+
response.raise_for_status()
23+
return await response.read()
24+
except aiohttp.ClientError as e:
25+
raise RuntimeError(f"Failed to download content from URL {url}: {str(e)}") from e

‎pyrit/converter/base_image_to_image_converter.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@
88
from typing import Literal
99
from urllib.parse import urlparse
1010

11-
import aiohttp
1211
from PIL import Image
1312

13+
from pyrit.converter._image_url_downloader import _download_image_from_url_async
1414
from pyrit.converter.converter import Converter, ConverterResult
1515
from pyrit.memory import data_serializer_factory
1616
from pyrit.models import PromptDataType
@@ -139,12 +139,7 @@ async def _read_image_from_url_async(self, url: str) -> bytes:
139139
Raises:
140140
RuntimeError: If there is an error during the download process.
141141
"""
142-
try:
143-
async with aiohttp.ClientSession() as session, session.get(url) as response:
144-
response.raise_for_status()
145-
return await response.read()
146-
except aiohttp.ClientError as e:
147-
raise RuntimeError(f"Failed to download content from URL {url}: {str(e)}") from e
142+
return await _download_image_from_url_async(url)
148143

149144
async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult:
150145
"""

‎pyrit/converter/image_compression_converter.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,9 @@
77
from typing import Any, Literal
88
from urllib.parse import urlparse
99

10-
import aiohttp
1110
from PIL import Image
1211

12+
from pyrit.converter._image_url_downloader import _download_image_from_url_async
1313
from pyrit.converter.converter import Converter, ConverterResult
1414
from pyrit.memory import data_serializer_factory
1515
from pyrit.models import ComponentIdentifier, PromptDataType
@@ -259,12 +259,7 @@ async def _read_image_from_url_async(self, url: str) -> bytes:
259259
Raises:
260260
RuntimeError: If there is an error during the download process.
261261
"""
262-
try:
263-
async with aiohttp.ClientSession() as session, session.get(url) as response:
264-
response.raise_for_status()
265-
return await response.read()
266-
except aiohttp.ClientError as e:
267-
raise RuntimeError(f"Failed to download content from URL {url}: {str(e)}") from e
262+
return await _download_image_from_url_async(url)
268263

269264
async def convert_async(self, *, prompt: str, input_type: PromptDataType = "image_path") -> ConverterResult:
270265
"""
Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT license.
3+
4+
import asyncio
5+
from io import BytesIO
6+
from unittest.mock import AsyncMock, MagicMock, patch
7+
8+
import aiohttp
9+
import pytest
10+
from PIL import Image
11+
12+
from pyrit.converter._image_url_downloader import _download_image_from_url_async
13+
from pyrit.converter.base_image_to_image_converter import BaseImageToImageConverter
14+
from pyrit.converter.image_compression_converter import ImageCompressionConverter
15+
16+
17+
class _OverridingImageConverter(BaseImageToImageConverter):
18+
"""Test converter that replaces URL reads through the protected hook."""
19+
20+
def __init__(self, *, image_bytes: bytes) -> None:
21+
super().__init__()
22+
self.image_bytes = image_bytes
23+
self.urls_read: list[str] = []
24+
25+
def _apply_transform(self, image: Image.Image) -> Image.Image:
26+
return image
27+
28+
async def _read_image_from_url_async(self, url: str) -> bytes:
29+
self.urls_read.append(url)
30+
return self.image_bytes
31+
32+
33+
def _mock_download_session() -> tuple[MagicMock, MagicMock, MagicMock]:
34+
client_session = MagicMock()
35+
session = MagicMock()
36+
response_context = MagicMock()
37+
response = MagicMock()
38+
39+
client_session.__aenter__ = AsyncMock(return_value=session)
40+
response_context.__aenter__ = AsyncMock(return_value=response)
41+
session.get.return_value = response_context
42+
43+
return client_session, session, response
44+
45+
46+
async def test_download_image_from_url_async_returns_response_bytes():
47+
"""The shared downloader returns raw response bytes after checking the status."""
48+
url = "https://example.com/image.png"
49+
client_session, session, response = _mock_download_session()
50+
response.read = AsyncMock(return_value=b"image bytes")
51+
52+
with patch("pyrit.converter._image_url_downloader.aiohttp.ClientSession", return_value=client_session):
53+
result = await _download_image_from_url_async(url)
54+
55+
assert result == b"image bytes"
56+
session.get.assert_called_once_with(url)
57+
response.raise_for_status.assert_called_once_with()
58+
response.read.assert_awaited_once_with()
59+
60+
61+
async def test_download_image_from_url_async_wraps_non_success_response():
62+
"""HTTP status failures retain the existing RuntimeError contract."""
63+
url = "https://example.com/missing.png"
64+
error = aiohttp.ClientResponseError(request_info=MagicMock(), history=(), status=404, message="Not Found")
65+
client_session, _, response = _mock_download_session()
66+
response.raise_for_status.side_effect = error
67+
68+
with (
69+
patch("pyrit.converter._image_url_downloader.aiohttp.ClientSession", return_value=client_session),
70+
pytest.raises(RuntimeError) as exc_info,
71+
):
72+
await _download_image_from_url_async(url)
73+
74+
assert type(exc_info.value) is RuntimeError
75+
assert str(exc_info.value) == f"Failed to download content from URL {url}: {str(error)}"
76+
assert exc_info.value.__cause__ is error
77+
78+
79+
async def test_download_image_from_url_async_wraps_connection_error():
80+
"""Connection failures retain the existing RuntimeError contract."""
81+
url = "https://example.com/image.png"
82+
error = aiohttp.ClientConnectionError("connection failed")
83+
client_session, session, _ = _mock_download_session()
84+
session.get.side_effect = error
85+
86+
with (
87+
patch("pyrit.converter._image_url_downloader.aiohttp.ClientSession", return_value=client_session),
88+
pytest.raises(RuntimeError) as exc_info,
89+
):
90+
await _download_image_from_url_async(url)
91+
92+
assert type(exc_info.value) is RuntimeError
93+
assert str(exc_info.value) == f"Failed to download content from URL {url}: {str(error)}"
94+
assert exc_info.value.__cause__ is error
95+
96+
97+
async def test_download_image_from_url_async_wraps_client_error_while_reading_body():
98+
"""Client errors while reading the response body retain the RuntimeError contract."""
99+
url = "https://example.com/image.png"
100+
error = aiohttp.ClientPayloadError("body read failed")
101+
client_session, _, response = _mock_download_session()
102+
response.read = AsyncMock(side_effect=error)
103+
104+
with (
105+
patch("pyrit.converter._image_url_downloader.aiohttp.ClientSession", return_value=client_session),
106+
pytest.raises(RuntimeError) as exc_info,
107+
):
108+
await _download_image_from_url_async(url)
109+
110+
assert type(exc_info.value) is RuntimeError
111+
assert str(exc_info.value) == f"Failed to download content from URL {url}: {str(error)}"
112+
assert exc_info.value.__cause__ is error
113+
114+
115+
async def test_download_image_from_url_async_propagates_non_client_body_error():
116+
"""Non-aiohttp body-read exceptions remain unwrapped."""
117+
url = "https://example.com/image.png"
118+
error = ValueError("body read failed")
119+
client_session, _, response = _mock_download_session()
120+
response.read = AsyncMock(side_effect=error)
121+
122+
with (
123+
patch("pyrit.converter._image_url_downloader.aiohttp.ClientSession", return_value=client_session),
124+
pytest.raises(ValueError) as exc_info,
125+
):
126+
await _download_image_from_url_async(url)
127+
128+
assert exc_info.value is error
129+
130+
131+
async def test_download_image_from_url_async_propagates_cancellation():
132+
"""Cancellation is not caught by the aiohttp error wrapper."""
133+
client_session, _, response = _mock_download_session()
134+
response.read = AsyncMock(side_effect=asyncio.CancelledError())
135+
136+
with (
137+
patch("pyrit.converter._image_url_downloader.aiohttp.ClientSession", return_value=client_session),
138+
pytest.raises(asyncio.CancelledError),
139+
):
140+
await _download_image_from_url_async("https://example.com/image.png")
141+
142+
143+
async def test_base_image_converter_url_reader_delegates_to_shared_downloader():
144+
"""The base hook delegates URL reads to the shared internal downloader."""
145+
converter = _OverridingImageConverter(image_bytes=b"")
146+
147+
with patch(
148+
"pyrit.converter.base_image_to_image_converter._download_image_from_url_async",
149+
new_callable=AsyncMock,
150+
return_value=b"image bytes",
151+
) as mock_download:
152+
result = await BaseImageToImageConverter._read_image_from_url_async(converter, "https://example.com/image.png")
153+
154+
assert result == b"image bytes"
155+
mock_download.assert_awaited_once_with("https://example.com/image.png")
156+
157+
158+
async def test_image_compression_converter_url_reader_delegates_to_shared_downloader():
159+
"""The compression hook delegates URL reads to the shared internal downloader."""
160+
converter = ImageCompressionConverter()
161+
162+
with patch(
163+
"pyrit.converter.image_compression_converter._download_image_from_url_async",
164+
new_callable=AsyncMock,
165+
return_value=b"image bytes",
166+
) as mock_download:
167+
result = await converter._read_image_from_url_async("https://example.com/image.png")
168+
169+
assert result == b"image bytes"
170+
mock_download.assert_awaited_once_with("https://example.com/image.png")
171+
172+
173+
async def test_base_image_converter_convert_async_honors_url_reader_override():
174+
"""URL conversion continues to dispatch through subclass URL-reader overrides."""
175+
image = Image.new("RGB", (1, 1), color=(0, 0, 0))
176+
image_buffer = BytesIO()
177+
image.save(image_buffer, format="PNG")
178+
converter = _OverridingImageConverter(image_bytes=image_buffer.getvalue())
179+
serializer = MagicMock()
180+
serializer.value = "converted.png"
181+
serializer.save_b64_image_async = AsyncMock()
182+
183+
with patch("pyrit.converter.base_image_to_image_converter.data_serializer_factory", return_value=serializer):
184+
result = await converter.convert_async(prompt="https://example.com/image.png", input_type="url")
185+
186+
assert result.output_text == "converted.png"
187+
assert converter.urls_read == ["https://example.com/image.png"]
188+
serializer.save_b64_image_async.assert_awaited_once()

0 commit comments

Comments
 (0)