Download feature and URLInputFile (#332)

* Fix How to upload docs

* Rename BaseBot to Bot

* Add download_file method

* Add download method

* Add URLInputFile

* Add Downloadable to __init__ and __all__

* Fix ImportError for Python 3.7

* Related pages

* Improving docs

* Some speed

* staticmethod to classmethod
This commit is contained in:
Gabben 2020-05-27 03:25:13 +05:00 committed by GitHub
parent 28382ebf5f
commit de3c5c1a8d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
40 changed files with 460 additions and 89 deletions

View file

@ -1,9 +1,22 @@
from __future__ import annotations
import datetime
import io
import pathlib
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, List, Optional, TypeVar, Union
from typing import (
Any,
AsyncGenerator,
AsyncIterator,
BinaryIO,
List,
Optional,
TypeVar,
Union,
cast,
)
import aiofiles
from async_lru import alru_cache
from ...utils.mixins import ContextInstanceMixin
@ -86,6 +99,7 @@ from ..types import (
Chat,
ChatMember,
ChatPermissions,
Downloadable,
File,
ForceReply,
GameHighScore,
@ -167,6 +181,93 @@ class Bot(ContextInstanceMixin["Bot"]):
"""
await self.session.close()
@classmethod
async def __download_file_binary_io(
cls, destination: BinaryIO, seek: bool, stream: AsyncGenerator[bytes, None]
) -> BinaryIO:
async for chunk in stream:
destination.write(chunk)
destination.flush()
if seek is True:
destination.seek(0)
return destination
@classmethod
async def __download_file(
cls, destination: Union[str, pathlib.Path], stream: AsyncGenerator[bytes, None]
) -> None:
async with aiofiles.open(destination, "wb") as f:
async for chunk in stream:
await f.write(chunk)
async def download_file(
self,
file_path: str,
destination: Optional[Union[BinaryIO, pathlib.Path, str]] = None,
timeout: int = 30,
chunk_size: int = 65536,
seek: bool = True,
) -> Optional[BinaryIO]:
"""
Download file by file_path to destination.
If you want to automatically create destination (:class:`io.BytesIO`) use default
value of destination and handle result of this method.
:param file_path: File path on Telegram server (You can get it from :obj:`aiogram.types.File`)
:param destination: Filename, file path or instance of :class:`io.IOBase`. For e.g. :class:`io.BytesIO`, defaults to None
:param timeout: Total timeout in seconds, defaults to 30
:param chunk_size: File chunks size, defaults to 64 kb
:param seek: Go to start of file when downloading is finished. Used only for destination with :class:`typing.BinaryIO` type, defaults to True
"""
if destination is None:
destination = io.BytesIO()
url = self.session.api.file_url(self.__token, file_path)
stream = self.session.stream_content(url=url, timeout=timeout, chunk_size=chunk_size)
if isinstance(destination, (str, pathlib.Path)):
return await self.__download_file(destination=destination, stream=stream)
else:
return await self.__download_file_binary_io(
destination=destination, seek=seek, stream=stream
)
async def download(
self,
file: Union[str, Downloadable],
destination: Optional[Union[BinaryIO, pathlib.Path, str]] = None,
timeout: int = 30,
chunk_size: int = 65536,
seek: bool = True,
) -> Optional[BinaryIO]:
"""
Download file by file_id or Downloadable object to destination.
If you want to automatically create destination (:class:`io.BytesIO`) use default
value of destination and handle result of this method.
:param file: file_id or Downloadable object
:param destination: Filename, file path or instance of :class:`io.IOBase`. For e.g. :class:`io.BytesIO`, defaults to None
:param timeout: Total timeout in seconds, defaults to 30
:param chunk_size: File chunks size, defaults to 64 kb
:param seek: Go to start of file when downloading is finished. Used only for destination with :class:`typing.BinaryIO` type, defaults to True
"""
if isinstance(file, str):
file_id = file
else:
file_id = getattr(file, "file_id", None)
if file_id is None:
raise TypeError("file can only be of the string or Downloadable type")
_file = await self.get_file(file_id)
# https://github.com/aiogram/aiogram/pull/282/files#r394110017
file_path = cast(str, _file.file_path)
return await self.download_file(
file_path, destination=destination, timeout=timeout, chunk_size=chunk_size, seek=seek
)
async def __call__(self, method: TelegramMethod[T]) -> T:
"""
Call API method

View file

@ -12,6 +12,7 @@ from .chosen_inline_result import ChosenInlineResult
from .contact import Contact
from .dice import Dice, DiceEmoji
from .document import Document
from .downloadable import Downloadable
from .encrypted_credentials import EncryptedCredentials
from .encrypted_passport_element import EncryptedPassportElement
from .file import File
@ -43,7 +44,7 @@ from .inline_query_result_venue import InlineQueryResultVenue
from .inline_query_result_video import InlineQueryResultVideo
from .inline_query_result_voice import InlineQueryResultVoice
from .input_contact_message_content import InputContactMessageContent
from .input_file import BufferedInputFile, FSInputFile, InputFile
from .input_file import BufferedInputFile, FSInputFile, InputFile, URLInputFile
from .input_location_message_content import InputLocationMessageContent
from .input_media import InputMedia
from .input_media_animation import InputMediaAnimation
@ -101,8 +102,10 @@ from .webhook_info import WebhookInfo
__all__ = (
"TelegramObject",
"Downloadable",
"BufferedInputFile",
"FSInputFile",
"URLInputFile",
"Update",
"WebhookInfo",
"User",

View file

@ -0,0 +1,5 @@
from typing_extensions import Protocol
class Downloadable(Protocol):
file_id: str

View file

@ -6,7 +6,7 @@ from abc import ABC, abstractmethod
from pathlib import Path
from typing import AsyncGenerator, AsyncIterator, Iterator, Optional, Union
import aiofiles as aiofiles
import aiofiles
DEFAULT_CHUNK_SIZE = 64 * 1024 # 64 kb
@ -82,3 +82,28 @@ class FSInputFile(InputFile):
while chunk:
yield chunk
chunk = await f.read(chunk_size)
class URLInputFile(InputFile):
def __init__(
self,
url: str,
filename: Optional[str] = None,
chunk_size: int = DEFAULT_CHUNK_SIZE,
timeout: int = 30,
):
super().__init__(filename=filename, chunk_size=chunk_size)
self.url = url
self.timeout = timeout
async def read(self, chunk_size: int) -> AsyncGenerator[bytes, None]:
from aiogram.api.client.bot import Bot
bot = Bot.get_current(no_error=False)
stream = bot.session.stream_content(
url=self.url, timeout=self.timeout, chunk_size=self.chunk_size
)
async for chunk in stream:
yield chunk

106
docs/api/download_file.md Normal file
View file

@ -0,0 +1,106 @@
# How to download file?
## Download file manually
First, you must get the `file_id` of the file you want to download. Information about files sent to the bot is contained in [Message](./types/message.md).
For example, download the document that came to the bot.
```python3
file_id = message.document.file_id
```
Then use the [getFile](./methods/get_file.md) method to get `file_path`.
```python3
file = await bot.get_file(file_id)
file_path = file.file_path
```
After that, use the [download_file](#download_file) method from the bot object.
### download_file(...)
Download file by `file_path` to destination.
If you want to automatically create destination (`#!python3 io.BytesIO`) use default
value of destination and handle result of this method.
|Argument|Type|Description|
|---|---|---|
| file_path | `#!python3 str` | File path on Telegram server |
| destination | `#!python3 Optional[Union[BinaryIO, pathlib.Path, str]]` | Filename, file path or instance of `#!python3 io.IOBase`. For e.g. `#!python3 io.BytesIO` (Default: `#!python3 None`) |
| timeout | `#!python3 int` | Total timeout in seconds (Default: `30`) |
| chunk_size | `#!python3 int` | File chunks size (Default: `64 kb`) |
| seek | `#!python3 bool` | Go to start of file when downloading is finished. Used only for destination with `#!python3 typing.BinaryIO` type (Default: `#!python3 True`) |
There are two options where you can download the file: to **disk** or to **binary I/O object**.
### Download file to disk
To download file to disk, you must specify the file name or path where to download the file. In this case, the function will return nothing.
```python3
await bot.download_file(file_path, "text.txt")
```
### Download file to binary I/O object
To download file to binary I/O object, you must specify an object with the `#!python3 typing.BinaryIO` type or use the default (`#!python3 None`) value.
In the first case, the function will return your object:
```python3
my_object = MyBinaryIO()
result: MyBinaryIO = await bot.download_file(file_path, my_object)
# print(result is my_object) # True
```
If you leave the default value, an `#!python3 io.BytesIO` object will be created and returned.
```python3
result: io.BytesIO = await bot.download_file(file_path)
```
## Download file in short way
Getting `file_path` manually every time is boring, so you should use the [download](#download) method.
### download(...)
Download file by `file_id` or `Downloadable` object to destination.
If you want to automatically create destination (`#!python3 io.BytesIO`) use default
value of destination and handle result of this method.
|Argument|Type|Description|
|---|---|---|
| file | `#!python3 Union[str, Downloadable]` | file_id or Downloadable object |
| destination | `#!python3 Optional[Union[BinaryIO, pathlib.Path, str]]` | Filename, file path or instance of `#!python3 io.IOBase`. For e.g. `#!python3 io.BytesIO` (Default: `#!python3 None`) |
| timeout | `#!python3 int` | Total timeout in seconds (Default: `30`) |
| chunk_size | `#!python3 int` | File chunks size (Default: `64 kb`) |
| seek | `#!python3 bool` | Go to start of file when downloading is finished. Used only for destination with `#!python3 typing.BinaryIO` type (Default: `#!python3 True`) |
It differs from [download_file](#download_file) **only** in that it accepts `file_id` or an `Downloadable` object (object that contains the `file_id` attribute) instead of `file_path`.
!!! note
All `Downloadable` objects are listed in Related pages.
You can download a file to [disk](#download-file-to-disk) or to a [binary I/O](#download-file-to-binary-io-object) object in the same way.
Example:
```python3
document = message.document
await bot.download(document)
```
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#getfile)
- [aiogram.types.Animation](types/animation.md)
- [aiogram.types.Audio](types/audio.md)
- [aiogram.types.Document](types/document.md)
- [aiogram.types.File](types/file.md)
- [aiogram.types.PassportFile](types/passport_file.md)
- [aiogram.types.PhotoSize](types/photo_size.md)
- [aiogram.types.Sticker](types/sticker.md)
- [aiogram.types.Video](types/video.md)
- [aiogram.types.VideoNote](types/video_note.md)
- [aiogram.types.Voice](types/voice.md)
- [How to upload file?](upload_file.md)

View file

@ -61,4 +61,4 @@ return AddStickerToSet(...)
- [Official documentation](https://core.telegram.org/bots/api#addstickertoset)
- [aiogram.types.InputFile](../types/input_file.md)
- [aiogram.types.MaskPosition](../types/mask_position.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -63,4 +63,4 @@ return CreateNewStickerSet(...)
- [Official documentation](https://core.telegram.org/bots/api#createnewstickerset)
- [aiogram.types.InputFile](../types/input_file.md)
- [aiogram.types.MaskPosition](../types/mask_position.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -70,4 +70,4 @@ return SendAnimation(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -72,4 +72,4 @@ return SendAudio(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -67,4 +67,4 @@ return SendDocument(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -66,4 +66,4 @@ return SendPhoto(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -64,4 +64,4 @@ return SendSticker(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -71,4 +71,4 @@ return SendVideo(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -67,4 +67,4 @@ return SendVideoNote(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -67,4 +67,4 @@ return SendVoice(...)
- [aiogram.types.Message](../types/message.md)
- [aiogram.types.ReplyKeyboardMarkup](../types/reply_keyboard_markup.md)
- [aiogram.types.ReplyKeyboardRemove](../types/reply_keyboard_remove.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -53,4 +53,4 @@ result: bool = await bot(SetChatPhoto(...))
- [Official documentation](https://core.telegram.org/bots/api#setchatphoto)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -57,4 +57,4 @@ return SetStickerSetThumb(...)
- [Official documentation](https://core.telegram.org/bots/api#setstickersetthumb)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -70,4 +70,4 @@ return SetWebhook(...)
- [Official documentation](https://core.telegram.org/bots/api#setwebhook)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -54,4 +54,4 @@ result: File = await bot(UploadStickerFile(...))
- [Official documentation](https://core.telegram.org/bots/api#uploadstickerfile)
- [aiogram.types.File](../types/file.md)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -31,3 +31,4 @@ This object represents an animation file (GIF or H.264/MPEG-4 AVC video without
- [Official documentation](https://core.telegram.org/bots/api#animation)
- [aiogram.types.PhotoSize](../types/photo_size.md)
- [How to download file?](../download_file.md)

View file

@ -30,3 +30,4 @@ This object represents an audio file to be treated as music by the Telegram clie
- [Official documentation](https://core.telegram.org/bots/api#audio)
- [aiogram.types.PhotoSize](../types/photo_size.md)
- [How to download file?](../download_file.md)

View file

@ -28,3 +28,4 @@ This object represents a general file (as opposed to photos, voice messages and
- [Official documentation](https://core.telegram.org/bots/api#document)
- [aiogram.types.PhotoSize](../types/photo_size.md)
- [How to download file?](../download_file.md)

View file

@ -27,3 +27,4 @@ Maximum file size to download is 20 MB
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#file)
- [How to download file?](../download_file.md)

View file

@ -16,4 +16,4 @@ This object represents the contents of a file to be uploaded. Must be posted usi
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#inputfile)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -30,4 +30,4 @@ Represents an animation file (GIF or H.264/MPEG-4 AVC video without sound) to be
- [Official documentation](https://core.telegram.org/bots/api#inputmediaanimation)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -30,4 +30,4 @@ Represents an audio file to be treated as music to be sent.
- [Official documentation](https://core.telegram.org/bots/api#inputmediaaudio)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -27,4 +27,4 @@ Represents a general file to be sent.
- [Official documentation](https://core.telegram.org/bots/api#inputmediadocument)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -26,4 +26,4 @@ Represents a photo to be sent.
- [Official documentation](https://core.telegram.org/bots/api#inputmediaphoto)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -31,4 +31,4 @@ Represents a video to be sent.
- [Official documentation](https://core.telegram.org/bots/api#inputmediavideo)
- [aiogram.types.InputFile](../types/input_file.md)
- [How to upload file?](../sending_files.md)
- [How to upload file?](../upload_file.md)

View file

@ -25,3 +25,4 @@ This object represents a file uploaded to Telegram Passport. Currently all Teleg
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#passportfile)
- [How to download file?](../download_file.md)

View file

@ -26,3 +26,4 @@ This object represents one size of a photo or a file / sticker thumbnail.
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#photosize)
- [How to download file?](../download_file.md)

View file

@ -33,3 +33,4 @@ This object represents a sticker.
- [Official documentation](https://core.telegram.org/bots/api#sticker)
- [aiogram.types.MaskPosition](../types/mask_position.md)
- [aiogram.types.PhotoSize](../types/photo_size.md)
- [How to download file?](../download_file.md)

View file

@ -30,3 +30,4 @@ This object represents a video file.
- [Official documentation](https://core.telegram.org/bots/api#video)
- [aiogram.types.PhotoSize](../types/photo_size.md)
- [How to download file?](../download_file.md)

View file

@ -28,3 +28,4 @@ This object represents a video message (available in Telegram apps as of v.4.0).
- [Official documentation](https://core.telegram.org/bots/api#videonote)
- [aiogram.types.PhotoSize](../types/photo_size.md)
- [How to download file?](../download_file.md)

View file

@ -26,3 +26,4 @@ This object represents a voice note.
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#voice)
- [How to download file?](../download_file.md)

View file

@ -3,10 +3,11 @@
As says [official Telegram Bot API documentation](https://core.telegram.org/bots/api#sending-files) there are three ways to send files (photos, stickers, audio, media, etc.):
If the file is already stored somewhere on the Telegram servers or file is available by the URL, you don't need to reupload it.
But if you need to upload new file just use subclasses of [InputFile](./types/input_file.md). Here is available two different types of input file:
But if you need to upload new file just use subclasses of [InputFile](./types/input_file.md). Here is available three different builtin types of input file:
- `#!python3 FSInputFile` - [uploading from file system](#upload-from-file-system)
- `#!python3 BufferedInputFile` - [uploading from buffer](#upload-from-buffer)
- `#!python3 URLInputFile` - [uploading from URL](#upload-from-url)
!!! warning "Be respectful with Telegram"
Instances of `InputFile` is reusable. That's mean you can create instance of InputFile and sent this file multiple times but Telegram is not recommend to do that and when you upload file once just save their `file_id` and use it in next times.
@ -65,3 +66,32 @@ file = BufferedInputFile.from_file("file.txt")
| `path` | `#!python3 Union[str, Path]` | File path |
| `filename` | `#!python3 Optional[str]` | Custom filename to be presented to Telegram |
| `chunk_size` | `#!python3 int` | File chunks size (Default: `64 kb`) |
## Upload from url
If you need to upload a file from another server, but the direct link is bound to your server's IP, or you want to bypass native [upload limits](https://core.telegram.org/bots/api#sending-files) by URL, you can use [URLInputFile](#urlinputfile).
Import wrapper:
```python3
from aiogram.types import URLInputFile
```
And then you can use it:
```python3
image = URLInputFile("https://www.python.org/static/community_logos/python-powered-h-140x182.png", filename="logo.png")
```
### URLInputFile(...)
|Argument|Type|Description|
|---|---|---|
| `url` | `#!python3 str` | URL |
| `filename` | `#!python3 Optional[str]` | Custom filename to be presented to Telegram |
| `chunk_size` | `#!python3 int` | File chunks size (Default: `64 kb`) |
| `timeout` | `#!python3 int` | Total timeout in seconds (Default: `30`) |
## Related pages:
- [Official documentation](https://core.telegram.org/bots/api#sending-files)
- [aiogram.types.InputFile](types/input_file.md)
- [How to download file?](download_file.md)

View file

@ -236,7 +236,8 @@ nav:
- api/types/game.md
- api/types/callback_game.md
- api/types/game_high_score.md
- api/sending_files.md
- api/download_file.md
- api/upload_file.md
- Dispatcher:
- dispatcher/index.md
- dispatcher/router.md

View file

@ -1,63 +0,0 @@
import pytest
from aiogram import Bot
from aiogram.api.client.session.aiohttp import AiohttpSession
from aiogram.api.methods import GetMe
try:
from asynctest import CoroutineMock, patch
except ImportError:
from unittest.mock import AsyncMock as CoroutineMock, patch # type: ignore
class TestBaseBot:
def test_init(self):
base_bot = Bot("42:TEST")
assert isinstance(base_bot.session, AiohttpSession)
assert base_bot.id == 42
def test_hashable(self):
base_bot = Bot("42:TEST")
assert hash(base_bot) == hash("42:TEST")
def test_equals(self):
base_bot = Bot("42:TEST")
assert base_bot == Bot("42:TEST")
assert base_bot != "42:TEST"
@pytest.mark.asyncio
async def test_emit(self):
base_bot = Bot("42:TEST")
method = GetMe()
with patch(
"aiogram.api.client.session.aiohttp.AiohttpSession.make_request",
new_callable=CoroutineMock,
) as mocked_make_request:
await base_bot(method)
mocked_make_request.assert_awaited_with("42:TEST", method)
@pytest.mark.asyncio
async def test_close(self):
base_bot = Bot("42:TEST", session=AiohttpSession())
await base_bot.session.create_session()
with patch(
"aiogram.api.client.session.aiohttp.AiohttpSession.close", new_callable=CoroutineMock
) as mocked_close:
await base_bot.close()
mocked_close.assert_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("close", [True, False])
async def test_context_manager(self, close: bool):
with patch(
"aiogram.api.client.session.aiohttp.AiohttpSession.close", new_callable=CoroutineMock
) as mocked_close:
async with Bot("42:TEST", session=AiohttpSession()).context(auto_close=close) as bot:
assert isinstance(bot, Bot)
if close:
mocked_close.assert_awaited()
else:
mocked_close.assert_not_awaited()

View file

@ -0,0 +1,133 @@
import io
import aiofiles
import pytest
from aresponses import ResponsesMockServer
from aiogram import Bot
from aiogram.api.client.session.aiohttp import AiohttpSession
from aiogram.api.methods import GetFile, GetMe
from aiogram.api.types import File, PhotoSize
from tests.mocked_bot import MockedBot
try:
from asynctest import CoroutineMock, patch
except ImportError:
from unittest.mock import AsyncMock as CoroutineMock, patch # type: ignore
class TestBot:
def test_init(self):
bot = Bot("42:TEST")
assert isinstance(bot.session, AiohttpSession)
assert bot.id == 42
def test_hashable(self):
bot = Bot("42:TEST")
assert hash(bot) == hash("42:TEST")
def test_equals(self):
bot = Bot("42:TEST")
assert bot == Bot("42:TEST")
assert bot != "42:TEST"
@pytest.mark.asyncio
async def test_emit(self):
bot = Bot("42:TEST")
method = GetMe()
with patch(
"aiogram.api.client.session.aiohttp.AiohttpSession.make_request",
new_callable=CoroutineMock,
) as mocked_make_request:
await bot(method)
mocked_make_request.assert_awaited_with("42:TEST", method)
@pytest.mark.asyncio
async def test_close(self):
bot = Bot("42:TEST", session=AiohttpSession())
await bot.session.create_session()
with patch(
"aiogram.api.client.session.aiohttp.AiohttpSession.close", new_callable=CoroutineMock
) as mocked_close:
await bot.close()
mocked_close.assert_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize("close", [True, False])
async def test_context_manager(self, close: bool):
with patch(
"aiogram.api.client.session.aiohttp.AiohttpSession.close", new_callable=CoroutineMock
) as mocked_close:
async with Bot("42:TEST", session=AiohttpSession()).context(auto_close=close) as bot:
assert isinstance(bot, Bot)
if close:
mocked_close.assert_awaited()
else:
mocked_close.assert_not_awaited()
@pytest.mark.asyncio
async def test_download_file(self, aresponses: ResponsesMockServer):
aresponses.add(
aresponses.ANY, aresponses.ANY, "get", aresponses.Response(status=200, body=b"\f" * 10)
)
# https://github.com/Tinche/aiofiles#writing-tests-for-aiofiles
aiofiles.threadpool.wrap.register(CoroutineMock)(
lambda *args, **kwargs: aiofiles.threadpool.AsyncBufferedIOBase(*args, **kwargs)
)
mock_file = CoroutineMock()
bot = Bot("42:TEST")
with patch("aiofiles.threadpool.sync_open", return_value=mock_file):
await bot.download_file("TEST", "file.png")
mock_file.write.assert_called_once_with(b"\f" * 10)
@pytest.mark.asyncio
async def test_download_file_default_destination(self, aresponses: ResponsesMockServer):
bot = Bot("42:TEST")
aresponses.add(
aresponses.ANY, aresponses.ANY, "get", aresponses.Response(status=200, body=b"\f" * 10)
)
result = await bot.download_file("TEST")
assert isinstance(result, io.BytesIO)
assert result.read() == b"\f" * 10
@pytest.mark.asyncio
async def test_download_file_custom_destination(self, aresponses: ResponsesMockServer):
bot = Bot("42:TEST")
aresponses.add(
aresponses.ANY, aresponses.ANY, "get", aresponses.Response(status=200, body=b"\f" * 10)
)
custom = io.BytesIO()
result = await bot.download_file("TEST", custom)
assert isinstance(result, io.BytesIO)
assert result is custom
assert result.read() == b"\f" * 10
@pytest.mark.asyncio
async def test_download(self, bot: MockedBot, aresponses: ResponsesMockServer):
bot.add_result_for(
GetFile, ok=True, result=File(file_id="file id", file_unique_id="file id")
)
bot.add_result_for(
GetFile, ok=True, result=File(file_id="file id", file_unique_id="file id")
)
assert await bot.download(File(file_id="file id", file_unique_id="file id"))
assert await bot.download("file id")
with pytest.raises(TypeError):
await bot.download(
[PhotoSize(file_id="file id", file_unique_id="file id", width=123, height=123)]
)

View file

@ -1,8 +1,10 @@
from typing import AsyncIterable
import pytest
from aresponses import ResponsesMockServer
from aiogram.api.types import BufferedInputFile, FSInputFile, InputFile
from aiogram import Bot
from aiogram.api.types import BufferedInputFile, FSInputFile, InputFile, URLInputFile
class TestInputFile:
@ -70,3 +72,21 @@ class TestInputFile:
assert chunk_size == 1
size += chunk_size
assert size > 0
@pytest.mark.asyncio
async def test_uri_input_file(self, aresponses: ResponsesMockServer):
aresponses.add(
aresponses.ANY, aresponses.ANY, "get", aresponses.Response(status=200, body=b"\f" * 10)
)
Bot.set_current(Bot("42:TEST"))
file = URLInputFile("https://test.org/", chunk_size=1)
size = 0
async for chunk in file:
assert chunk == b"\f"
chunk_size = len(chunk)
assert chunk_size == 1
size += chunk_size
assert size == 10