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