diff --git a/deepinfra/utils/read_stream.py b/deepinfra/utils/read_stream.py index 1c2a3c1..81de154 100644 --- a/deepinfra/utils/read_stream.py +++ b/deepinfra/utils/read_stream.py @@ -3,6 +3,14 @@ import httpx +from deepinfra._exceptions import ( + APIConnectionError, + APIStatusError, + APITimeoutError, +) + +DOWNLOAD_TIMEOUT = 30.0 + class ReadStreamUtils: """ @@ -34,8 +42,24 @@ def url_to_stream(url): Downloads an image from a URL and returns it as a BytesIO. :param url: The URL of the image. :return: A BytesIO containing the image data. + :raises APIStatusError: the URL answered with an error status. Without + this the error page's bytes would be uploaded as the file. + :raises APIConnectionError: the download never got a response. """ - response = httpx.get(url, follow_redirects=True) + try: + response = httpx.get( + url, follow_redirects=True, timeout=DOWNLOAD_TIMEOUT + ) + except httpx.TimeoutException as exc: + raise APITimeoutError(f"Timed out downloading {url}") from exc + except httpx.TransportError as exc: + raise APIConnectionError(f"Failed to download {url}: {exc}") from exc + if response.is_error: + raise APIStatusError( + f"Failed to download {url}", + status_code=response.status_code, + response=response, + ) return BytesIO(response.content) @staticmethod diff --git a/tests/test_read_stream.py b/tests/test_read_stream.py new file mode 100644 index 0000000..39cc064 --- /dev/null +++ b/tests/test_read_stream.py @@ -0,0 +1,73 @@ +"""URL inputs must fail loudly. + +A blob field given as a URL used to be fetched without checking the response, +so a 404 page was uploaded to the API as if it were the file. +""" + +import base64 + +import httpx +import pytest +import respx + +from deepinfra import APIConnectionError, APIStatusError, APITimeoutError +from deepinfra.utils.form_data import FormDataUtils +from deepinfra.utils.read_stream import ReadStreamUtils + +URL = "https://example.com/audio.mp3" + + +@respx.mock +def test_url_download_returns_content(): + respx.get(URL).mock(return_value=httpx.Response(200, content=b"audio-bytes")) + + assert ReadStreamUtils.get_read_stream(URL).read() == b"audio-bytes" + + +@respx.mock +@pytest.mark.parametrize("status", [404, 401, 500]) +def test_error_status_raises_instead_of_uploading_the_error_page(status): + respx.get(URL).mock(return_value=httpx.Response(status, html="