From 7ab0db79919659615d4291f8958396f1ff1c0829 Mon Sep 17 00:00:00 2001 From: gabbhack <43146729+gabbhack@users.noreply.github.com> Date: Wed, 22 Jan 2020 22:54:29 +0500 Subject: [PATCH 1/2] Add stream_content to BaseSession and add tests --- aiogram/api/client/session/base.py | 8 ++++++- tests/mocked_bot.py | 7 ++++++- .../test_session/test_base_session.py | 21 ++++++++++++++++++- 3 files changed, 33 insertions(+), 3 deletions(-) diff --git a/aiogram/api/client/session/base.py b/aiogram/api/client/session/base.py index f2f40727..f31dd451 100644 --- a/aiogram/api/client/session/base.py +++ b/aiogram/api/client/session/base.py @@ -3,7 +3,7 @@ from __future__ import annotations import abc import datetime import json -from typing import Any, Callable, Optional, TypeVar, Union +from typing import Any, AsyncGenerator, Callable, Optional, TypeVar, Union from aiogram.utils.exceptions import TelegramAPIError @@ -44,6 +44,12 @@ class BaseSession(abc.ABC): async def make_request(self, token: str, method: TelegramMethod[T]) -> T: # pragma: no cover pass + @abc.abstractmethod + async def stream_content( + self, url: str, timeout: int, chunk_size: int + ) -> AsyncGenerator[bytes, None]: # pragma: no cover + yield b"" + def prepare_value(self, value: Any) -> Union[str, int, bool]: if isinstance(value, str): return value diff --git a/tests/mocked_bot.py b/tests/mocked_bot.py index 60d16d52..9ee5731c 100644 --- a/tests/mocked_bot.py +++ b/tests/mocked_bot.py @@ -1,5 +1,5 @@ from collections import deque -from typing import TYPE_CHECKING, Deque, Optional, Type +from typing import TYPE_CHECKING, AsyncGenerator, Deque, Optional, Type from aiogram import Bot from aiogram.api.client.session.base import BaseSession @@ -29,6 +29,11 @@ class MockedSession(BaseSession): self.raise_for_status(response) return response.result # type: ignore + async def stream_content( + self, url: str, timeout: int, chunk_size: int + ) -> AsyncGenerator[bytes, None]: # pragma: no cover + yield b"" + class MockedBot(Bot): if TYPE_CHECKING: diff --git a/tests/test_api/test_client/test_session/test_base_session.py b/tests/test_api/test_client/test_session/test_base_session.py index f5004b68..cc565a42 100644 --- a/tests/test_api/test_client/test_session/test_base_session.py +++ b/tests/test_api/test_client/test_session/test_base_session.py @@ -1,5 +1,5 @@ import datetime -from typing import AsyncContextManager +from typing import AsyncContextManager, AsyncGenerator import pytest @@ -22,6 +22,14 @@ class CustomSession(BaseSession): assert isinstance(token, str) assert isinstance(method, TelegramMethod) + async def stream_content( + self, url: str, timeout: int, chunk_size: int + ) -> AsyncGenerator[bytes, None]: # pragma: no cover + assert isinstance(url, str) + assert isinstance(timeout, int) + assert isinstance(chunk_size, int) + yield b"\f" * 10 + class TestBaseSession(DataMixin): def test_init_api(self): @@ -100,6 +108,17 @@ class TestBaseSession(DataMixin): assert await session.make_request("42:TEST", GetMe()) is None + @pytest.mark.asyncio + async def test_stream_content(self): + session = CustomSession() + stream = session.stream_content( + "https://www.python.org/static/img/python-logo.png", timeout=5, chunk_size=65536 + ) + assert isinstance(stream, AsyncGenerator) + + async for chunk in stream: + assert isinstance(chunk, bytes) + @pytest.mark.asyncio async def test_context_manager(self): session = CustomSession() From 26708154b0c131d96795ad594e56f2a210656d8a Mon Sep 17 00:00:00 2001 From: gabbhack <43146729+gabbhack@users.noreply.github.com> Date: Wed, 22 Jan 2020 22:55:34 +0500 Subject: [PATCH 2/2] Implement stream_content in AiohttpSession and add tests --- aiogram/api/client/session/aiohttp.py | 14 ++++++++++-- .../test_session/test_aiohttp_session.py | 22 ++++++++++++++++++- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/aiogram/api/client/session/aiohttp.py b/aiogram/api/client/session/aiohttp.py index 9ac73eaa..06a1c99d 100644 --- a/aiogram/api/client/session/aiohttp.py +++ b/aiogram/api/client/session/aiohttp.py @@ -1,8 +1,8 @@ from __future__ import annotations -from typing import Callable, Optional, TypeVar, cast +from typing import AsyncGenerator, Callable, Optional, TypeVar, cast -from aiohttp import ClientSession, FormData +from aiohttp import ClientSession, ClientTimeout, FormData from aiogram.api.methods import Request, TelegramMethod @@ -56,6 +56,16 @@ class AiohttpSession(BaseSession): self.raise_for_status(response) return cast(T, response.result) + async def stream_content( + self, url: str, timeout: int, chunk_size: int + ) -> AsyncGenerator[bytes, None]: + session = await self.create_session() + client_timeout = ClientTimeout(total=timeout) + + async with session.get(url, timeout=client_timeout) as resp: + async for chunk in resp.content.iter_chunked(chunk_size): + yield chunk + async def __aenter__(self) -> AiohttpSession: await self.create_session() return self 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 edde8057..ec1e16a5 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 @@ -1,4 +1,4 @@ -from typing import AsyncContextManager +from typing import AsyncContextManager, AsyncGenerator import aiohttp import pytest @@ -107,6 +107,26 @@ class TestAiohttpSession: assert patched_raise_for_status.called_once() + @pytest.mark.asyncio + async def test_stream_content(self, aresponses: ResponsesMockServer): + aresponses.add( + aresponses.ANY, aresponses.ANY, "get", aresponses.Response(status=200, body=b"\f" * 10) + ) + + session = AiohttpSession() + stream = session.stream_content( + "https://www.python.org/static/img/python-logo.png", timeout=5, chunk_size=1 + ) + assert isinstance(stream, AsyncGenerator) + + size = 0 + async for chunk in stream: + assert isinstance(chunk, bytes) + chunk_size = len(chunk) + assert chunk_size == 1 + size += chunk_size + assert size == 10 + @pytest.mark.asyncio async def test_context_manager(self): session = AiohttpSession()