mirror of
https://github.com/aiogram/aiogram.git
synced 2025-12-12 18:19:34 +00:00
[3.x] Check status code when downloading file (#1079)
* Check status code when downloading file and raise an error if someting bad happends * Style fixes * Add doc * Use "towncrier create <issue>.<type>" for creating file
This commit is contained in:
parent
94e11ce8e9
commit
184ee1fbf8
7 changed files with 45 additions and 7 deletions
1
CHANGES/816.bugfix
Normal file
1
CHANGES/816.bugfix
Normal file
|
|
@ -0,0 +1 @@
|
||||||
|
Check status code when downloading file
|
||||||
|
|
@ -304,7 +304,12 @@ class Bot(ContextInstanceMixin["Bot"]):
|
||||||
close_stream = True
|
close_stream = True
|
||||||
else:
|
else:
|
||||||
url = self.session.api.file_url(self.__token, file_path)
|
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:
|
try:
|
||||||
if isinstance(destination, (str, pathlib.Path)):
|
if isinstance(destination, (str, pathlib.Path)):
|
||||||
|
|
|
||||||
|
|
@ -162,11 +162,11 @@ class AiohttpSession(BaseSession):
|
||||||
return cast(TelegramType, response.result)
|
return cast(TelegramType, response.result)
|
||||||
|
|
||||||
async def stream_content(
|
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]:
|
) -> AsyncGenerator[bytes, None]:
|
||||||
session = await self.create_session()
|
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):
|
async for chunk in resp.content.iter_chunked(chunk_size):
|
||||||
yield chunk
|
yield chunk
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -153,7 +153,7 @@ class BaseSession(abc.ABC):
|
||||||
|
|
||||||
@abc.abstractmethod
|
@abc.abstractmethod
|
||||||
async def stream_content(
|
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
|
) -> AsyncGenerator[bytes, None]: # pragma: no cover
|
||||||
"""
|
"""
|
||||||
Stream reader
|
Stream reader
|
||||||
|
|
|
||||||
|
|
@ -135,7 +135,10 @@ class URLInputFile(InputFile):
|
||||||
|
|
||||||
bot = Bot.get_current(no_error=False)
|
bot = Bot.get_current(no_error=False)
|
||||||
stream = bot.session.stream_content(
|
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:
|
async for chunk in stream:
|
||||||
|
|
|
||||||
|
|
@ -37,7 +37,11 @@ class MockedSession(BaseSession):
|
||||||
return response.result # type: ignore
|
return response.result # type: ignore
|
||||||
|
|
||||||
async def stream_content(
|
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
|
) -> AsyncGenerator[bytes, None]: # pragma: no cover
|
||||||
yield b""
|
yield b""
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -190,7 +190,10 @@ class TestAiohttpSession:
|
||||||
|
|
||||||
session = AiohttpSession()
|
session = AiohttpSession()
|
||||||
stream = session.stream_content(
|
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)
|
assert isinstance(stream, AsyncGenerator)
|
||||||
|
|
||||||
|
|
@ -202,6 +205,28 @@ class TestAiohttpSession:
|
||||||
size += chunk_size
|
size += chunk_size
|
||||||
assert size == 10
|
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):
|
async def test_context_manager(self):
|
||||||
session = AiohttpSession()
|
session = AiohttpSession()
|
||||||
assert isinstance(session, AsyncContextManager)
|
assert isinstance(session, AsyncContextManager)
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue