Add stream_content to BaseSession and add tests

This commit is contained in:
gabbhack 2020-01-22 22:54:29 +05:00
parent 213f2a3c16
commit 7ab0db7991
3 changed files with 33 additions and 3 deletions

View file

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