Test InputFile

This commit is contained in:
Alex Root Junior 2019-11-20 00:32:19 +02:00
parent fa6732542d
commit 9bc1aae0f5
2 changed files with 77 additions and 6 deletions

View file

@ -4,7 +4,7 @@ import io
import os
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional, Union
from typing import AsyncGenerator, Optional, Union
import aiofiles as aiofiles
@ -28,8 +28,8 @@ class InputFile(ABC):
yield
@abstractmethod
async def read(self, chunk_size: int):
pass
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
yield b""
async def __aiter__(self):
async for chunk in self.read(self.chunk_size):
@ -48,14 +48,14 @@ class BufferedInputFile(InputFile):
path: Union[str, Path],
filename: Optional[str] = None,
chunk_size: int = DEFAULT_CHUNK_SIZE,
):
) -> BufferedInputFile:
if filename is None:
filename = os.path.basename(path)
with open(path, "rb") as f:
data = f.read()
return cls(data, filename=filename, chunk_size=chunk_size)
async def read(self, chunk_size: int):
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
buffer = io.BytesIO(self.data)
chunk = buffer.read(chunk_size)
while chunk:
@ -76,7 +76,7 @@ class FSInputFile(InputFile):
self.path = path
async def read(self, chunk_size: int):
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
async with aiofiles.open(self.path, "rb") as f:
chunk = await f.read(chunk_size)
while chunk:

View file

@ -0,0 +1,71 @@
from typing import AsyncIterable
import pytest
from aiogram.api.types import BufferedInputFile, FSInputFile, InputFile
class TestInputFile:
def test_fs_input_file(self):
file = FSInputFile(__file__)
assert isinstance(file, InputFile)
assert isinstance(file, AsyncIterable)
assert file.filename is not None
assert file.filename.startswith("test_")
assert file.filename.endswith(".py")
assert file.chunk_size > 0
@pytest.mark.asyncio
async def test_fs_input_file_readable(self):
file = FSInputFile(__file__, chunk_size=1)
assert file.chunk_size == 1
size = 0
async for chunk in file:
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
assert size > 0
def test_buffered_input_file(self):
file = BufferedInputFile(b"\f" * 10, filename="file.bin")
assert isinstance(file, InputFile)
assert isinstance(file, AsyncIterable)
assert file.filename == "file.bin"
assert isinstance(file.data, bytes)
@pytest.mark.asyncio
async def test_buffered_input_file_readable(self):
file = BufferedInputFile(b"\f" * 10, filename="file.bin", chunk_size=1)
size = 0
async for chunk in file:
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
assert size == 10
@pytest.mark.asyncio
async def test_buffered_input_file_from_file(self):
file = BufferedInputFile.from_file(__file__, chunk_size=10)
assert isinstance(file, InputFile)
assert isinstance(file, AsyncIterable)
assert file.filename is not None
assert file.filename.startswith("test_")
assert file.filename.endswith(".py")
assert isinstance(file.data, bytes)
assert file.chunk_size == 10
@pytest.mark.asyncio
async def test_buffered_input_file_from_file_readable(self):
file = BufferedInputFile.from_file(__file__, chunk_size=1)
size = 0
async for chunk in file:
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
assert size > 0