diff --git a/CHANGES/816.bugfix b/CHANGES/816.bugfix new file mode 100644 index 00000000..22bfee8c --- /dev/null +++ b/CHANGES/816.bugfix @@ -0,0 +1 @@ +Check status code when downloading file diff --git a/aiogram/client/bot.py b/aiogram/client/bot.py index fc1198e3..d9ddc07a 100644 --- a/aiogram/client/bot.py +++ b/aiogram/client/bot.py @@ -304,7 +304,12 @@ class Bot(ContextInstanceMixin["Bot"]): close_stream = True else: url = self.session.api.file_url(self.__token, file_path) - stream = self.session.stream_content(url=url, timeout=timeout, chunk_size=chunk_size) + stream = self.session.stream_content( + url=url, + timeout=timeout, + chunk_size=chunk_size, + raise_for_status=True, + ) try: if isinstance(destination, (str, pathlib.Path)): diff --git a/aiogram/client/session/aiohttp.py b/aiogram/client/session/aiohttp.py index 66aae752..4d9dff06 100644 --- a/aiogram/client/session/aiohttp.py +++ b/aiogram/client/session/aiohttp.py @@ -162,11 +162,11 @@ class AiohttpSession(BaseSession): return cast(TelegramType, response.result) async def stream_content( - self, url: str, timeout: int, chunk_size: int + self, url: str, timeout: int, chunk_size: int, raise_for_status: bool ) -> AsyncGenerator[bytes, None]: session = await self.create_session() - async with session.get(url, timeout=timeout) as resp: + async with session.get(url, timeout=timeout, raise_for_status=raise_for_status) as resp: async for chunk in resp.content.iter_chunked(chunk_size): yield chunk diff --git a/aiogram/client/session/base.py b/aiogram/client/session/base.py index 95392139..971d9a32 100644 --- a/aiogram/client/session/base.py +++ b/aiogram/client/session/base.py @@ -153,7 +153,7 @@ class BaseSession(abc.ABC): @abc.abstractmethod async def stream_content( - self, url: str, timeout: int, chunk_size: int + self, url: str, timeout: int, chunk_size: int, raise_for_status: bool ) -> AsyncGenerator[bytes, None]: # pragma: no cover """ Stream reader diff --git a/aiogram/types/input_file.py b/aiogram/types/input_file.py index 102e365b..c4090505 100644 --- a/aiogram/types/input_file.py +++ b/aiogram/types/input_file.py @@ -135,7 +135,10 @@ class URLInputFile(InputFile): bot = Bot.get_current(no_error=False) stream = bot.session.stream_content( - url=self.url, timeout=self.timeout, chunk_size=self.chunk_size + url=self.url, + timeout=self.timeout, + chunk_size=self.chunk_size, + raise_for_status=True, ) async for chunk in stream: diff --git a/tests/mocked_bot.py b/tests/mocked_bot.py index 1783eb86..25a9b38e 100644 --- a/tests/mocked_bot.py +++ b/tests/mocked_bot.py @@ -37,7 +37,11 @@ class MockedSession(BaseSession): return response.result # type: ignore async def stream_content( - self, url: str, timeout: int, chunk_size: int + self, + url: str, + timeout: int, + chunk_size: int, + raise_for_status: bool, ) -> AsyncGenerator[bytes, None]: # pragma: no cover yield b"" diff --git a/tests/test_api/test_client/test_session/test_aiohttp_session.py b/tests/test_api/test_client/test_session/test_aiohttp_session.py index 14c25ad3..651711cb 100644 --- a/tests/test_api/test_client/test_session/test_aiohttp_session.py +++ b/tests/test_api/test_client/test_session/test_aiohttp_session.py @@ -190,7 +190,10 @@ class TestAiohttpSession: session = AiohttpSession() stream = session.stream_content( - "https://www.python.org/static/img/python-logo.png", timeout=5, chunk_size=1 + "https://www.python.org/static/img/python-logo.png", + timeout=5, + chunk_size=1, + raise_for_status=True, ) assert isinstance(stream, AsyncGenerator) @@ -202,6 +205,28 @@ class TestAiohttpSession: size += chunk_size assert size == 10 + async def test_stream_content_404(self, aresponses: ResponsesMockServer): + aresponses.add( + aresponses.ANY, + aresponses.ANY, + "get", + aresponses.Response( + status=404, + body=b"File not found", + ), + ) + session = AiohttpSession() + stream = session.stream_content( + "https://www.python.org/static/img/python-logo.png", + timeout=5, + chunk_size=1, + raise_for_status=True, + ) + + with pytest.raises(ClientError): + async for _ in stream: + ... + async def test_context_manager(self): session = AiohttpSession() assert isinstance(session, AsyncContextManager)