mirror of
https://github.com/aiogram/aiogram.git
synced 2025-12-06 07:50:32 +00:00
Test InputFile
This commit is contained in:
parent
fa6732542d
commit
9bc1aae0f5
2 changed files with 77 additions and 6 deletions
|
|
@ -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:
|
||||
|
|
|
|||
71
tests/test_api/test_types/test_input_file.py
Normal file
71
tests/test_api/test_types/test_input_file.py
Normal 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
|
||||
Loading…
Add table
Add a link
Reference in a new issue