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()