Move prepare_<field> to single module and add tests for default parse_mode

This commit is contained in:
Alex Root Junior 2019-12-26 01:35:21 +02:00
parent df6a861276
commit e12aca9046
22 changed files with 165 additions and 100 deletions

View file

@ -1,7 +1,7 @@
from typing import Any, Dict, Optional, Union from typing import Any, Dict, Optional, Union
from ..types import InputFile, MaskPosition from ..types import InputFile, MaskPosition
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class AddStickerToSet(TelegramMethod[bool]): class AddStickerToSet(TelegramMethod[bool]):
@ -31,6 +31,6 @@ class AddStickerToSet(TelegramMethod[bool]):
data: Dict[str, Any] = self.dict(exclude={"png_sticker"}) data: Dict[str, Any] = self.dict(exclude={"png_sticker"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="png_sticker", value=self.png_sticker) prepare_file(data=data, files=files, name="png_sticker", value=self.png_sticker)
return Request(method="addStickerToSet", data=data, files=files) return Request(method="addStickerToSet", data=data, files=files)

View file

@ -1,7 +1,7 @@
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from ..types import InlineQueryResult from ..types import InlineQueryResult
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_parse_mode
class AnswerInlineQuery(TelegramMethod[bool]): class AnswerInlineQuery(TelegramMethod[bool]):
@ -38,6 +38,6 @@ class AnswerInlineQuery(TelegramMethod[bool]):
def build_request(self) -> Request: def build_request(self) -> Request:
data: Dict[str, Any] = self.dict() data: Dict[str, Any] = self.dict()
self.prepare_parse_mode(data["results"]) prepare_parse_mode(data["results"])
return Request(method="answerInlineQuery", data=data) return Request(method="answerInlineQuery", data=data)

View file

@ -2,7 +2,7 @@ from __future__ import annotations
import abc import abc
import secrets import secrets
from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar, Union
from pydantic import BaseConfig, BaseModel, Extra from pydantic import BaseConfig, BaseModel, Extra
from pydantic.generics import GenericModel from pydantic.generics import GenericModel
@ -53,7 +53,17 @@ class TelegramMethod(abc.ABC, BaseModel, Generic[T]):
# noinspection PyTypeChecker # noinspection PyTypeChecker
return Response[self.__returning__](**data) # type: ignore return Response[self.__returning__](**data) # type: ignore
def prepare_file(self, name: str, value: Any, data: Dict[str, Any], files: Dict[str, Any]): async def emit(self, bot: Bot) -> T:
return await bot.emit(self)
def __await__(self):
from aiogram.api.client.bot import Bot
bot = Bot.get_current(no_error=False)
return self.emit(bot).__await__()
def prepare_file(name: str, value: Any, data: Dict[str, Any], files: Dict[str, Any]):
if not value: if not value:
return return
if name == "thumb": if name == "thumb":
@ -65,10 +75,34 @@ class TelegramMethod(abc.ABC, BaseModel, Generic[T]):
else: else:
data[name] = value data[name] = value
def prepare_parse_mode(self, root: Any) -> None:
def prepare_input_media(data: Dict[str, Any], files: Dict[str, InputFile]) -> None:
for input_media in data.get("media", []): # type: Dict[str, Union[str, InputFile]]
if (
"media" in input_media
and input_media["media"]
and isinstance(input_media["media"], InputFile)
):
tag = secrets.token_urlsafe(10)
files[tag] = input_media.pop("media") # type: ignore
input_media["media"] = f"attach://{tag}"
def prepare_media_file(data: Dict[str, Any], files: Dict[str, InputFile]) -> None:
if (
data["media"]
and "media" in data["media"]
and isinstance(data["media"]["media"], InputFile)
):
tag = secrets.token_urlsafe(10)
files[tag] = data["media"].pop("media") # type: ignore
data["media"]["media"] = f"attach://{tag}"
def prepare_parse_mode(root: Any) -> None:
if isinstance(root, list): if isinstance(root, list):
for item in root: for item in root:
self.prepare_parse_mode(item) prepare_parse_mode(item)
return return
if root.get("parse_mode"): if root.get("parse_mode"):
@ -81,12 +115,3 @@ class TelegramMethod(abc.ABC, BaseModel, Generic[T]):
root["parse_mode"] = bot.parse_mode root["parse_mode"] = bot.parse_mode
return return
return return
async def emit(self, bot: Bot) -> T:
return await bot.emit(self)
def __await__(self):
from aiogram.api.client.bot import Bot
bot = Bot.get_current(no_error=False)
return self.emit(bot).__await__()

View file

@ -1,7 +1,7 @@
from typing import Any, Dict, Optional, Union from typing import Any, Dict, Optional, Union
from ..types import InputFile, MaskPosition from ..types import InputFile, MaskPosition
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class CreateNewStickerSet(TelegramMethod[bool]): class CreateNewStickerSet(TelegramMethod[bool]):
@ -39,6 +39,6 @@ class CreateNewStickerSet(TelegramMethod[bool]):
data: Dict[str, Any] = self.dict(exclude={"png_sticker"}) data: Dict[str, Any] = self.dict(exclude={"png_sticker"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="png_sticker", value=self.png_sticker) prepare_file(data=data, files=files, name="png_sticker", value=self.png_sticker)
return Request(method="createNewStickerSet", data=data, files=files) return Request(method="createNewStickerSet", data=data, files=files)

View file

@ -1,7 +1,7 @@
from typing import Any, Dict, Optional, Union from typing import Any, Dict, Optional, Union
from ..types import InlineKeyboardMarkup, Message from ..types import InlineKeyboardMarkup, Message
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_parse_mode
class EditMessageCaption(TelegramMethod[Union[Message, bool]]): class EditMessageCaption(TelegramMethod[Union[Message, bool]]):
@ -31,6 +31,6 @@ class EditMessageCaption(TelegramMethod[Union[Message, bool]]):
def build_request(self) -> Request: def build_request(self) -> Request:
data: Dict[str, Any] = self.dict() data: Dict[str, Any] = self.dict()
self.prepare_parse_mode(data) prepare_parse_mode(data)
return Request(method="editMessageCaption", data=data) return Request(method="editMessageCaption", data=data)

View file

@ -1,8 +1,7 @@
import secrets
from typing import Any, Dict, Optional, Union from typing import Any, Dict, Optional, Union
from ..types import InlineKeyboardMarkup, InputFile, InputMedia, Message from ..types import InlineKeyboardMarkup, InputFile, InputMedia, Message
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_media_file, prepare_parse_mode
class EditMessageMedia(TelegramMethod[Union[Message, bool]]): class EditMessageMedia(TelegramMethod[Union[Message, bool]]):
@ -33,19 +32,9 @@ class EditMessageMedia(TelegramMethod[Union[Message, bool]]):
def build_request(self) -> Request: def build_request(self) -> Request:
data: Dict[str, Any] = self.dict() data: Dict[str, Any] = self.dict()
self.prepare_parse_mode(data["media"]) prepare_parse_mode(data["media"])
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_media_file(data=data, files=files) prepare_media_file(data=data, files=files)
return Request(method="editMessageMedia", data=data, files=files) return Request(method="editMessageMedia", data=data, files=files)
def prepare_media_file(self, data: Dict[str, Any], files: Dict[str, InputFile]) -> None:
if (
data["media"]
and "media" in data["media"]
and isinstance(data["media"]["media"], InputFile)
):
tag = secrets.token_urlsafe(10)
files[tag] = data["media"].pop("media") # type: ignore
data["media"]["media"] = f"attach://{tag}"

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendAnimation(TelegramMethod[Message]): class SendAnimation(TelegramMethod[Message]):
@ -61,7 +61,7 @@ class SendAnimation(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"animation", "thumb"}) data: Dict[str, Any] = self.dict(exclude={"animation", "thumb"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="animation", value=self.animation) prepare_file(data=data, files=files, name="animation", value=self.animation)
self.prepare_file(data=data, files=files, name="thumb", value=self.thumb) prepare_file(data=data, files=files, name="thumb", value=self.thumb)
return Request(method="sendAnimation", data=data, files=files) return Request(method="sendAnimation", data=data, files=files)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendAudio(TelegramMethod[Message]): class SendAudio(TelegramMethod[Message]):
@ -63,7 +63,7 @@ class SendAudio(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"audio", "thumb"}) data: Dict[str, Any] = self.dict(exclude={"audio", "thumb"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="audio", value=self.audio) prepare_file(data=data, files=files, name="audio", value=self.audio)
self.prepare_file(data=data, files=files, name="thumb", value=self.thumb) prepare_file(data=data, files=files, name="thumb", value=self.thumb)
return Request(method="sendAudio", data=data, files=files) return Request(method="sendAudio", data=data, files=files)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendDocument(TelegramMethod[Message]): class SendDocument(TelegramMethod[Message]):
@ -55,7 +55,7 @@ class SendDocument(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"document", "thumb"}) data: Dict[str, Any] = self.dict(exclude={"document", "thumb"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="document", value=self.document) prepare_file(data=data, files=files, name="document", value=self.document)
self.prepare_file(data=data, files=files, name="thumb", value=self.thumb) prepare_file(data=data, files=files, name="thumb", value=self.thumb)
return Request(method="sendDocument", data=data, files=files) return Request(method="sendDocument", data=data, files=files)

View file

@ -1,8 +1,7 @@
import secrets
from typing import Any, Dict, List, Optional, Union from typing import Any, Dict, List, Optional, Union
from ..types import InputFile, InputMediaPhoto, InputMediaVideo, Message from ..types import InputFile, InputMediaPhoto, InputMediaVideo, Message
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_input_media, prepare_parse_mode
class SendMediaGroup(TelegramMethod[List[Message]]): class SendMediaGroup(TelegramMethod[List[Message]]):
@ -27,21 +26,9 @@ class SendMediaGroup(TelegramMethod[List[Message]]):
def build_request(self) -> Request: def build_request(self) -> Request:
data: Dict[str, Any] = self.dict() data: Dict[str, Any] = self.dict()
self.prepare_parse_mode(data["media"]) prepare_parse_mode(data["media"])
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_input_media(data, files) prepare_input_media(data, files)
return Request(method="sendMediaGroup", data=data, files=files) return Request(method="sendMediaGroup", data=data, files=files)
@staticmethod
def prepare_input_media(data: Dict[str, Any], files: Dict[str, InputFile]) -> None:
for input_media in data.get("media", []): # type: Dict[str, Union[str, InputFile]]
if (
"media" in input_media
and input_media["media"]
and isinstance(input_media["media"], InputFile)
):
tag = secrets.token_urlsafe(10)
files[tag] = input_media.pop("media") # type: ignore
input_media["media"] = f"attach://{tag}"

View file

@ -7,7 +7,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_parse_mode
class SendMessage(TelegramMethod[Message]): class SendMessage(TelegramMethod[Message]):
@ -41,6 +41,6 @@ class SendMessage(TelegramMethod[Message]):
def build_request(self) -> Request: def build_request(self) -> Request:
data: Dict[str, Any] = self.dict() data: Dict[str, Any] = self.dict()
self.prepare_parse_mode(data) prepare_parse_mode(data)
return Request(method="sendMessage", data=data) return Request(method="sendMessage", data=data)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendPhoto(TelegramMethod[Message]): class SendPhoto(TelegramMethod[Message]):
@ -46,6 +46,6 @@ class SendPhoto(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"photo"}) data: Dict[str, Any] = self.dict(exclude={"photo"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="photo", value=self.photo) prepare_file(data=data, files=files, name="photo", value=self.photo)
return Request(method="sendPhoto", data=data, files=files) return Request(method="sendPhoto", data=data, files=files)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendSticker(TelegramMethod[Message]): class SendSticker(TelegramMethod[Message]):
@ -42,6 +42,6 @@ class SendSticker(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"sticker"}) data: Dict[str, Any] = self.dict(exclude={"sticker"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="sticker", value=self.sticker) prepare_file(data=data, files=files, name="sticker", value=self.sticker)
return Request(method="sendSticker", data=data, files=files) return Request(method="sendSticker", data=data, files=files)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendVideo(TelegramMethod[Message]): class SendVideo(TelegramMethod[Message]):
@ -63,7 +63,7 @@ class SendVideo(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"video", "thumb"}) data: Dict[str, Any] = self.dict(exclude={"video", "thumb"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="video", value=self.video) prepare_file(data=data, files=files, name="video", value=self.video)
self.prepare_file(data=data, files=files, name="thumb", value=self.thumb) prepare_file(data=data, files=files, name="thumb", value=self.thumb)
return Request(method="sendVideo", data=data, files=files) return Request(method="sendVideo", data=data, files=files)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendVideoNote(TelegramMethod[Message]): class SendVideoNote(TelegramMethod[Message]):
@ -53,7 +53,7 @@ class SendVideoNote(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"video_note", "thumb"}) data: Dict[str, Any] = self.dict(exclude={"video_note", "thumb"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="video_note", value=self.video_note) prepare_file(data=data, files=files, name="video_note", value=self.video_note)
self.prepare_file(data=data, files=files, name="thumb", value=self.thumb) prepare_file(data=data, files=files, name="thumb", value=self.thumb)
return Request(method="sendVideoNote", data=data, files=files) return Request(method="sendVideoNote", data=data, files=files)

View file

@ -8,7 +8,7 @@ from ..types import (
ReplyKeyboardMarkup, ReplyKeyboardMarkup,
ReplyKeyboardRemove, ReplyKeyboardRemove,
) )
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SendVoice(TelegramMethod[Message]): class SendVoice(TelegramMethod[Message]):
@ -52,6 +52,6 @@ class SendVoice(TelegramMethod[Message]):
data: Dict[str, Any] = self.dict(exclude={"voice"}) data: Dict[str, Any] = self.dict(exclude={"voice"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="voice", value=self.voice) prepare_file(data=data, files=files, name="voice", value=self.voice)
return Request(method="sendVoice", data=data, files=files) return Request(method="sendVoice", data=data, files=files)

View file

@ -1,7 +1,7 @@
from typing import Any, Dict, Union from typing import Any, Dict, Union
from ..types import InputFile from ..types import InputFile
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SetChatPhoto(TelegramMethod[bool]): class SetChatPhoto(TelegramMethod[bool]):
@ -25,6 +25,6 @@ class SetChatPhoto(TelegramMethod[bool]):
data: Dict[str, Any] = self.dict(exclude={"photo"}) data: Dict[str, Any] = self.dict(exclude={"photo"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="photo", value=self.photo) prepare_file(data=data, files=files, name="photo", value=self.photo)
return Request(method="setChatPhoto", data=data, files=files) return Request(method="setChatPhoto", data=data, files=files)

View file

@ -1,7 +1,7 @@
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
from ..types import InputFile from ..types import InputFile
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class SetWebhook(TelegramMethod[bool]): class SetWebhook(TelegramMethod[bool]):
@ -46,6 +46,6 @@ class SetWebhook(TelegramMethod[bool]):
data: Dict[str, Any] = self.dict(exclude={"certificate"}) data: Dict[str, Any] = self.dict(exclude={"certificate"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="certificate", value=self.certificate) prepare_file(data=data, files=files, name="certificate", value=self.certificate)
return Request(method="setWebhook", data=data, files=files) return Request(method="setWebhook", data=data, files=files)

View file

@ -1,7 +1,7 @@
from typing import Any, Dict from typing import Any, Dict
from ..types import File, InputFile from ..types import File, InputFile
from .base import Request, TelegramMethod from .base import Request, TelegramMethod, prepare_file
class UploadStickerFile(TelegramMethod[File]): class UploadStickerFile(TelegramMethod[File]):
@ -24,6 +24,6 @@ class UploadStickerFile(TelegramMethod[File]):
data: Dict[str, Any] = self.dict(exclude={"png_sticker"}) data: Dict[str, Any] = self.dict(exclude={"png_sticker"})
files: Dict[str, InputFile] = {} files: Dict[str, InputFile] = {}
self.prepare_file(data=data, files=files, name="png_sticker", value=self.png_sticker) prepare_file(data=data, files=files, name="png_sticker", value=self.png_sticker)
return Request(method="uploadStickerFile", data=data, files=files) return Request(method="uploadStickerFile", data=data, files=files)

View file

@ -34,8 +34,8 @@ class MockedBot(Bot):
if TYPE_CHECKING: if TYPE_CHECKING:
session: MockedSession session: MockedSession
def __init__(self): def __init__(self, **kwargs):
super(MockedBot, self).__init__("42:TEST", session=MockedSession()) super(MockedBot, self).__init__("42:TEST", session=MockedSession(), **kwargs)
def add_result_for( def add_result_for(
self, self,

View file

@ -0,0 +1,62 @@
from typing import Dict, Optional
import pytest
from aiogram import Bot
from aiogram.api.methods.base import prepare_parse_mode
class TestPrepareFile:
# TODO: Add tests
pass
class TestPrepareInputMedia:
# TODO: Add tests
pass
class TestPrepareMediaFile:
# TODO: Add tests
pass
class TestPrepareParseMode:
@pytest.mark.parametrize(
"parse_mode,data,result",
[
[None, {}, None],
["HTML", {}, "HTML"],
["Markdown", {}, "Markdown"],
[None, {"parse_mode": "HTML"}, "HTML"],
["HTML", {"parse_mode": "HTML"}, "HTML"],
["Markdown", {"parse_mode": "HTML"}, "HTML"],
],
)
@pytest.mark.asyncio
async def test_default_parse_mode(
self, parse_mode: str, data: Dict[str, str], result: Optional[str]
):
async with Bot(token="42:TEST", parse_mode=parse_mode).context() as bot:
assert bot.parse_mode == parse_mode
prepare_parse_mode(data)
assert data.get("parse_mode") == result
@pytest.mark.asyncio
async def test_list(self):
data = [{}] * 2
data.append({"parse_mode": "HTML"})
async with Bot(token="42:TEST", parse_mode="Markdown").context():
prepare_parse_mode(data)
assert isinstance(data, list)
assert len(data) == 3
assert all("parse_mode" in item for item in data)
assert data[0]["parse_mode"] == "Markdown"
assert data[1]["parse_mode"] == "Markdown"
assert data[2]["parse_mode"] == "HTML"
def test_bot_not_in_context(self):
data = {}
prepare_parse_mode(data)
assert "parse_mode" not in data

View file

@ -191,6 +191,8 @@ class TestDispatcher:
def test_run_polling(self, bot: MockedBot): def test_run_polling(self, bot: MockedBot):
dispatcher = Dispatcher() dispatcher = Dispatcher()
with patch("aiogram.dispatcher.dispatcher.Dispatcher.start_polling") as patched_start_polling: with patch(
"aiogram.dispatcher.dispatcher.Dispatcher.start_polling"
) as patched_start_polling:
dispatcher.run_polling(bot) dispatcher.run_polling(bot)
patched_start_polling.assert_awaited_once() patched_start_polling.assert_awaited_once()