├── src └── aiogram_mock │ ├── __init__.py │ ├── py.typed │ ├── facade_factory.py │ ├── tg_control.py │ ├── mocked_session.py │ └── tg_state.py ├── .flake8 ├── mypy.ini ├── .gitignore ├── requirements.txt ├── Makefile ├── .github └── workflows │ ├── publish-to-pypi.yml │ └── test-package.yml ├── LICENSE ├── examples └── test_basic.py ├── pyproject.toml └── README.md /src/aiogram_mock/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/aiogram_mock/py.typed: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.flake8: -------------------------------------------------------------------------------- 1 | [flake8] 2 | max-line-length=120 3 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | plugins = pydantic.mypy 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea/ 2 | build/ 3 | dist/ 4 | env/ 5 | vitya.egg-info/ 6 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # CI 2 | pytest==7.3.1 3 | pytest-asyncio 4 | 5 | mypy==1.2.0 6 | flake8==6.0.0 7 | isort==5.12.0 8 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | .PHONY: lint-flake8 2 | 3 | lint-flake8: 4 | @echo run flake8 5 | @flake8 src/ 6 | 7 | .PHONY: lint-isort 8 | 9 | lint-isort: 10 | @echo run isort 11 | @isort src/ 12 | 13 | .PHONY: lint-mypy 14 | 15 | lint-mypy: 16 | @echo run mypy 17 | @mypy src 18 | 19 | lint: lint-flake8 lint-isort lint-mypy 20 | -------------------------------------------------------------------------------- /.github/workflows/publish-to-pypi.yml: -------------------------------------------------------------------------------- 1 | name: Publish package to PyPI 2 | 3 | on: 4 | release: 5 | types: [ released ] 6 | 7 | jobs: 8 | build-n-publish: 9 | name: Build and publish package to PyPI 10 | runs-on: ubuntu-latest 11 | steps: 12 | - uses: actions/checkout@master 13 | - name: Set up Python 3.8 14 | uses: actions/setup-python@v1 15 | with: 16 | python-version: 3.8 17 | 18 | - name: Install pypa/build 19 | run: >- 20 | python3 -m pip install --user --upgrade setuptools wheel 21 | - name: Build a binary wheel and a source tarball 22 | run: >- 23 | python3 setup.py sdist bdist_wheel 24 | 25 | - name: Publish package to PyPI 26 | if: startsWith(github.ref, 'refs/tags') 27 | uses: pypa/gh-action-pypi-publish@master 28 | with: 29 | password: ${{ secrets.PYPI_API_TOKEN }} 30 | -------------------------------------------------------------------------------- /.github/workflows/test-package.yml: -------------------------------------------------------------------------------- 1 | name: Python package 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | branches: [ master ] 8 | 9 | jobs: 10 | build: 11 | runs-on: ubuntu-latest 12 | strategy: 13 | matrix: 14 | python-version: [3.8, 3.9, 3.10.11, 3.11.2] 15 | 16 | steps: 17 | - uses: actions/checkout@master 18 | - name: Set up Python ${{ matrix.python-version }} 19 | uses: actions/setup-python@master 20 | with: 21 | python-version: ${{ matrix.python-version }} 22 | 23 | - name: Install dependencies 24 | run: | 25 | python -m pip install --upgrade pip 26 | pip install -r requirements.txt 27 | pip install -e . 28 | 29 | - name: Lint with flake8 30 | run: | 31 | make lint-flake8 32 | 33 | - name: Check isort 34 | run: | 35 | make lint-isort 36 | 37 | - name: Test with mypy 38 | run: | 39 | make lint-mypy 40 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 HICEBANK 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /examples/test_basic.py: -------------------------------------------------------------------------------- 1 | from typing import Tuple, Generator 2 | 3 | import pytest 4 | from aiogram import Bot, Dispatcher, F 5 | from aiogram.filters import CommandStart 6 | from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, CallbackQuery 7 | 8 | from aiogram_mock.facade_factory import private_chat_tg_control 9 | from aiogram_mock.tg_control import PrivateChatTgControl 10 | 11 | 12 | async def on_start(message: Message): 13 | await message.answer( 14 | 'hello', 15 | reply_markup=InlineKeyboardMarkup( 16 | inline_keyboard=[ 17 | [ 18 | InlineKeyboardButton(text='ping', callback_data='ping') 19 | ] 20 | ] 21 | ) 22 | ) 23 | 24 | 25 | async def on_click_me(query: CallbackQuery): 26 | await query.answer(text='pong') 27 | 28 | 29 | def create_bot_and_dispatcher() -> Tuple[Bot, Dispatcher]: 30 | bot = Bot(token='123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11') 31 | dispatcher = Dispatcher() 32 | dispatcher.message.register(on_start, CommandStart()) 33 | dispatcher.callback_query.register(on_click_me) 34 | return bot, dispatcher 35 | 36 | 37 | @pytest.fixture() 38 | def tg_control() -> Generator[PrivateChatTgControl, None, None]: 39 | bot, dispatcher = create_bot_and_dispatcher() 40 | with private_chat_tg_control( 41 | bot=bot, 42 | dispatcher=dispatcher, 43 | ) as tg_control: 44 | yield tg_control 45 | 46 | 47 | async def test_start(tg_control): 48 | await tg_control.send("/start") 49 | assert tg_control.last_message.text == 'hello' 50 | 51 | answer = await tg_control.click(F.text == 'ping') 52 | assert answer.text == 'pong' 53 | 54 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ['setuptools==67.7.1'] 3 | build-backend = 'setuptools.build_meta' 4 | 5 | [project] 6 | name = 'aiogram_mock' 7 | version = '0.0.1' 8 | description = 'Tools for testing of aiogram applications' 9 | readme = 'README.md' 10 | requires-python = '>=3.8' 11 | 12 | classifiers = [ 13 | 'Development Status :: 3 - Alpha', 14 | 'Programming Language :: Python :: 3', 15 | 'Programming Language :: Python :: 3.8', 16 | 'Programming Language :: Python :: 3.9', 17 | 'Programming Language :: Python :: 3.10', 18 | 'Programming Language :: Python :: Implementation :: CPython', 19 | 'Programming Language :: Python :: Implementation :: PyPy', 20 | 'Operating System :: OS Independent', 21 | 'Intended Audience :: Developers', 22 | 'Intended Audience :: Information Technology', 23 | 'License :: OSI Approved :: Apache Software License', 24 | 'Typing :: Typed', 25 | 'Topic :: Software Development :: Libraries :: Python Modules', 26 | 'Topic :: Internet', 27 | ] 28 | 29 | dependencies = [ 30 | 'aiogram>=3.0.0b6' 31 | ] 32 | 33 | [project.urls] 34 | 'Homepage' = 'https://github.com/hicebank/aiogram_mock' 35 | 'Bug Tracker' = 'https://github.com/hicebank/aiogram_mock/issues' 36 | 'Source' = 'https://github.com/hicebank/aiogram_mock/' 37 | 'Download' = 'https://github.com/hicebank/aiogram_mock/#files' 38 | 39 | [tool.setuptools] 40 | package-dir = {"" = "src"} 41 | 42 | [tool.setuptools.package-data] 43 | adaptix = ['py.typed'] 44 | 45 | [[project.authors]] 46 | name = 'Pavel' 47 | 48 | 49 | [tool.pytest.ini_options] 50 | asyncio_mode = 'auto' 51 | 52 | [tool.isort] 53 | multi_line_output = 3 54 | line_length = 120 55 | include_trailing_comma = true 56 | combine_as_imports = true 57 | remove_redundant_aliases = true 58 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # aiogram_mock 2 | 3 | Tools for testing of aiogram applications 4 | 5 | ``` 6 | pip install git+https://github.com/hicebank/aiogram-mock#egg=aiogram_mock 7 | ``` 8 | 9 | 10 | Example 11 | ```python 12 | from typing import Tuple, Generator 13 | 14 | import pytest 15 | from aiogram import Bot, Dispatcher, F 16 | from aiogram.filters import CommandStart 17 | from aiogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton, CallbackQuery 18 | 19 | from aiogram_mock.facade_factory import private_chat_tg_control 20 | from aiogram_mock.tg_control import PrivateChatTgControl 21 | 22 | 23 | async def on_start(message: Message): 24 | await message.answer( 25 | 'hello', 26 | reply_markup=InlineKeyboardMarkup( 27 | inline_keyboard=[ 28 | [ 29 | InlineKeyboardButton(text='ping', callback_data='ping') 30 | ] 31 | ] 32 | ) 33 | ) 34 | 35 | 36 | async def on_click_me(query: CallbackQuery): 37 | await query.answer(text='pong') 38 | 39 | 40 | def create_bot_and_dispatcher() -> Tuple[Bot, Dispatcher]: 41 | bot = Bot(token='123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11') 42 | dispatcher = Dispatcher() 43 | dispatcher.message.register(on_start, CommandStart()) 44 | dispatcher.callback_query.register(on_click_me) 45 | return bot, dispatcher 46 | 47 | 48 | @pytest.fixture() 49 | def tg_control() -> Generator[PrivateChatTgControl, None, None]: 50 | bot, dispatcher = create_bot_and_dispatcher() 51 | with private_chat_tg_control( 52 | bot=bot, 53 | dispatcher=dispatcher, 54 | ) as tg_control: 55 | yield tg_control 56 | 57 | 58 | async def test_start(tg_control): 59 | await tg_control.send("/start") 60 | assert tg_control.last_message.text == 'hello' 61 | 62 | answer = await tg_control.click(F.text == 'ping') 63 | assert answer.text == 'pong' 64 | ``` 65 | 66 | -------------------------------------------------------------------------------- /src/aiogram_mock/facade_factory.py: -------------------------------------------------------------------------------- 1 | from contextlib import contextmanager 2 | from typing import Callable, Generator, Iterable, Optional 3 | from unittest.mock import patch 4 | 5 | from aiogram import Bot, Dispatcher 6 | from aiogram.client.session.base import BaseSession 7 | from aiogram.types import Chat, User 8 | 9 | from aiogram_mock.mocked_session import MockedSession 10 | from aiogram_mock.tg_control import PrivateChatTgControl, TgControl 11 | from aiogram_mock.tg_state import TgState 12 | 13 | 14 | @contextmanager 15 | def private_chat_tg_control( 16 | dispatcher: Dispatcher, 17 | bot: Bot, 18 | target_user: Optional[User] = None, 19 | bot_user: Optional[User] = None, 20 | tg_state_factory: Callable[[Iterable[Chat]], TgState] = TgState, 21 | mocked_session_factory: Callable[[TgState, User], BaseSession] = MockedSession, 22 | ) -> Generator[PrivateChatTgControl, None, None]: 23 | if target_user is None: 24 | target_user = User( 25 | id=103592704, 26 | first_name='Linus', 27 | last_name='Torvalds', 28 | is_bot=False, 29 | ) 30 | if bot_user is None: 31 | bot_user = User( 32 | id=738453453, 33 | first_name='Test', 34 | last_name='bot', 35 | username='test_bot', 36 | is_bot=True, 37 | ) 38 | 39 | chat = Chat( 40 | id=target_user.id, 41 | type='private', 42 | username=target_user.username, 43 | first_name=target_user.first_name, 44 | last_name=target_user.last_name, 45 | ) 46 | 47 | tg_state = tg_state_factory([chat]) 48 | session = mocked_session_factory(tg_state, bot_user) 49 | with patch.object(bot, 'session', session): 50 | yield PrivateChatTgControl( 51 | tg_control=TgControl( 52 | dispatcher=dispatcher, 53 | bot=bot, 54 | tg_state=tg_state, 55 | ), 56 | chat=chat, 57 | user=target_user, 58 | ) 59 | -------------------------------------------------------------------------------- /src/aiogram_mock/tg_control.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | from datetime import datetime 3 | from typing import Callable, Optional, Sequence, Union 4 | 5 | from aiogram import Bot, Dispatcher, MagicFilter 6 | from aiogram.fsm.context import FSMContext 7 | from aiogram.fsm.storage.base import DEFAULT_DESTINY, BaseStorage, StorageKey 8 | from aiogram.methods import AnswerCallbackQuery 9 | from aiogram.types import CallbackQuery, Chat, Contact, InlineKeyboardButton, Message, Update, User 10 | 11 | from aiogram_mock.tg_state import TgState, UserState 12 | 13 | 14 | class TgControl: 15 | def __init__( 16 | self, 17 | dispatcher: Dispatcher, 18 | bot: Bot, 19 | tg_state: TgState, 20 | ): 21 | self._dispatcher = dispatcher 22 | self._bot = bot 23 | self._tg_state = tg_state 24 | 25 | def messages(self, chat_id: int) -> Sequence[Message]: 26 | return self._tg_state.chat_history(chat_id) 27 | 28 | def last_message(self, chat_id: int) -> Message: 29 | return self._tg_state.chat_history(chat_id)[-1] 30 | 31 | def user_state(self, *, chat_id: int, user_id: int) -> UserState: 32 | return self._tg_state.get_user_state(chat_id=chat_id, user_id=user_id) 33 | 34 | async def _send_message(self, message: Message) -> None: 35 | await self._dispatcher.feed_update( 36 | self._bot, 37 | Update( 38 | update_id=self._tg_state.increment_update_id(), 39 | message=self._tg_state.add_message(message), 40 | ), 41 | ) 42 | 43 | async def send(self, from_user: User, chat: Chat, text: str) -> None: 44 | await self._send_message( 45 | Message( 46 | message_id=self._tg_state.next_message_id(chat.id), 47 | date=datetime.utcnow(), 48 | from_user=from_user, 49 | chat=chat, 50 | text=text, 51 | ), 52 | ) 53 | 54 | async def send_contact(self, from_user: User, chat: Chat, contact: Contact) -> None: 55 | await self._send_message( 56 | Message( 57 | message_id=self._tg_state.next_message_id(chat.id), 58 | date=datetime.utcnow(), 59 | from_user=from_user, 60 | chat=chat, 61 | contact=contact, 62 | ), 63 | ) 64 | 65 | async def click( 66 | self, 67 | selector: Union[Callable[[InlineKeyboardButton], bool], MagicFilter], 68 | message: Message, 69 | user: User, 70 | ) -> AnswerCallbackQuery: 71 | if message.reply_markup is None: 72 | raise ValueError('Message has no inline keyboard') 73 | if isinstance(selector, MagicFilter): 74 | selector = selector.resolve 75 | 76 | buttons = itertools.chain.from_iterable(message.reply_markup.inline_keyboard) 77 | selected_buttons = [button for button in buttons if selector(button)] 78 | if len(selected_buttons) == 0: 79 | raise ValueError('selector skip all buttons') 80 | if len(selected_buttons) > 1: 81 | raise ValueError('selector selects more than one button') 82 | button = selected_buttons[0] 83 | 84 | callback_query_id = self._tg_state.next_callback_query_id() 85 | await self._dispatcher.feed_update( 86 | self._bot, 87 | Update( 88 | update_id=self._tg_state.increment_update_id(), 89 | callback_query=CallbackQuery( 90 | id=callback_query_id, 91 | data=button.callback_data, 92 | chat_instance=str(message.chat.id), # message.chat.id contains local id 93 | from_user=user, 94 | message=message, 95 | ), 96 | ), 97 | ) 98 | return self._tg_state.get_answer_callback_query(callback_query_id) 99 | 100 | @property 101 | def bot(self) -> Bot: 102 | return self._bot 103 | 104 | @property 105 | def storage(self) -> BaseStorage: 106 | return self._dispatcher.storage 107 | 108 | 109 | class PrivateChatTgControl: 110 | def __init__(self, tg_control: TgControl, chat: Chat, user: User): 111 | self._tg_control = tg_control 112 | self._chat = chat 113 | self._user = user 114 | self._validate() 115 | 116 | def _validate(self) -> None: 117 | if self._chat.id != self._user.id: 118 | raise ValueError('chat.id and user.id must be equal') 119 | 120 | def state(self, destiny: str = DEFAULT_DESTINY) -> FSMContext: 121 | return FSMContext( 122 | bot=self.bot, 123 | storage=self._tg_control.storage, 124 | key=StorageKey( 125 | bot_id=self.bot.id, 126 | chat_id=self._chat.id, 127 | user_id=self._user.id, 128 | destiny=destiny, 129 | ), 130 | ) 131 | 132 | @property 133 | def messages(self) -> Sequence[Message]: 134 | return self._tg_control.messages(self._chat.id) 135 | 136 | @property 137 | def last_message(self) -> Message: 138 | return self._tg_control.last_message(self._chat.id) 139 | 140 | @property 141 | def user_state(self) -> UserState: 142 | return self._tg_control.user_state(chat_id=self._chat.id, user_id=self._user.id) 143 | 144 | @property 145 | def bot(self) -> Bot: 146 | return self._tg_control.bot 147 | 148 | @property 149 | def user(self) -> User: 150 | return self._user 151 | 152 | @property 153 | def chat(self) -> Chat: 154 | return self._chat 155 | 156 | async def send(self, text: str) -> None: 157 | return await self._tg_control.send( 158 | from_user=self._user, 159 | chat=self._chat, 160 | text=text, 161 | ) 162 | 163 | async def send_contact(self, contact: Contact) -> None: 164 | await self._tg_control.send_contact( 165 | from_user=self._user, 166 | chat=self._chat, 167 | contact=contact, 168 | ) 169 | 170 | async def click( 171 | self, 172 | selector: Union[Callable[[InlineKeyboardButton], bool], MagicFilter], 173 | message: Optional[Message] = None, 174 | ) -> AnswerCallbackQuery: 175 | if message is None: 176 | message = self.last_message 177 | return await self._tg_control.click(selector, message, self._user) 178 | -------------------------------------------------------------------------------- /src/aiogram_mock/mocked_session.py: -------------------------------------------------------------------------------- 1 | from datetime import datetime 2 | from typing import Any, AsyncGenerator, List, Mapping, Optional, Sequence, Type, Union 3 | 4 | from aiogram import Bot 5 | from aiogram.client.session.base import BaseSession 6 | from aiogram.methods import ( 7 | AnswerCallbackQuery, 8 | EditMessageReplyMarkup, 9 | EditMessageText, 10 | SendMessage, 11 | SendPhoto, 12 | SetChatMenuButton, 13 | TelegramMethod, 14 | ) 15 | from aiogram.methods.base import TelegramType 16 | from aiogram.types import UNSET, InlineKeyboardMarkup, Message, ReplyKeyboardRemove, User 17 | 18 | from aiogram_mock.tg_state import TgState 19 | 20 | SendMessageVariant = Union[SendMessage, SendPhoto] 21 | 22 | 23 | class MockedSession(BaseSession): 24 | def __init__(self, tg_state: TgState, bot_user: User): 25 | self._tg_state = tg_state 26 | self._bot_user = bot_user 27 | self._sent_methods: List[TelegramMethod[Any]] = [] 28 | super().__init__() 29 | 30 | async def close(self) -> None: 31 | pass 32 | 33 | def _process_reply_markup(self, chat_id: int, method: SendMessageVariant) -> Optional[InlineKeyboardMarkup]: 34 | if isinstance(method.reply_markup, InlineKeyboardMarkup): 35 | return method.reply_markup 36 | elif isinstance(method.reply_markup, ReplyKeyboardRemove): 37 | self._tg_state.update_chat_user_state(chat_id=chat_id, reply_markup=None) 38 | return None 39 | if method.reply_markup is not None: 40 | self._tg_state.update_chat_user_state(chat_id=chat_id, reply_markup=method.reply_markup) 41 | return None 42 | 43 | def _process_reply(self, chat_id: int, method: SendMessageVariant) -> Optional[Message]: 44 | if method.reply_to_message_id is not None: 45 | try: 46 | return self._tg_state.get_message(chat_id, method.reply_to_message_id) 47 | except KeyError: 48 | if not method.allow_sending_without_reply: 49 | raise 50 | return None 51 | 52 | async def _mock_send_message(self, bot: Bot, method: SendMessage, timeout: Optional[int] = UNSET) -> Message: 53 | chat_id = int(method.chat_id) 54 | return self._tg_state.add_message( 55 | Message( 56 | message_id=self._tg_state.next_message_id(chat_id), 57 | text=method.text, 58 | chat=self._tg_state.chats[chat_id], 59 | date=datetime.utcnow(), 60 | message_thread_id=method.message_thread_id, 61 | from_user=self._bot_user, 62 | reply_to_message=self._process_reply(chat_id, method), 63 | reply_markup=self._process_reply_markup(chat_id, method), 64 | ), 65 | ) 66 | 67 | async def _mock_send_photo(self, bot: Bot, method: SendPhoto, timeout: Optional[int] = UNSET) -> Message: 68 | chat_id = int(method.chat_id) 69 | return self._tg_state.add_message( 70 | Message( 71 | message_id=self._tg_state.next_message_id(chat_id), 72 | text=method.caption, 73 | chat=self._tg_state.chats[chat_id], 74 | date=datetime.utcnow(), 75 | message_thread_id=method.message_thread_id, 76 | from_user=self._bot_user, 77 | reply_to_message=self._process_reply(chat_id, method), 78 | reply_markup=self._process_reply_markup(chat_id, method), 79 | document=await self._tg_state.load_file(bot.id, method.photo), 80 | ), 81 | ) 82 | 83 | async def _mock_answer_callback_query( 84 | self, 85 | bot: Bot, 86 | method: AnswerCallbackQuery, 87 | timeout: Optional[Message] = UNSET, 88 | ) -> bool: 89 | self._tg_state.add_answer_callback_query(method) 90 | return True 91 | 92 | async def _mock_set_chat_menu_button( 93 | self, 94 | bot: Bot, 95 | method: AnswerCallbackQuery, 96 | timeout: Optional[int] = UNSET, 97 | ) -> bool: 98 | return True 99 | 100 | async def _mock_edit_message_text( 101 | self, 102 | bot: Bot, 103 | method: EditMessageText, 104 | timeout: Optional[int] = UNSET, 105 | ) -> Union[Message, bool]: 106 | if method.chat_id is None and method.message_id is None: 107 | raise NotImplementedError('Editing of inlined message is not supported') 108 | if method.chat_id is None or method.message_id is None: 109 | raise ValueError('Bad sent message') 110 | 111 | message = self._tg_state.get_message(int(method.chat_id), method.message_id) 112 | new_message = message.copy( 113 | update={'text': method.text, 'reply_markup': method.reply_markup}, 114 | ) 115 | self._tg_state.replace_message(new_message) 116 | return new_message 117 | 118 | async def _mock_edit_message_reply_markup( 119 | self, 120 | bot: Bot, 121 | method: EditMessageReplyMarkup, 122 | timeout: Optional[int] = UNSET, 123 | ) -> Union[Message, bool]: 124 | if method.chat_id is None and method.message_id is None: 125 | raise NotImplementedError('Edititng of inlined message is not supported') 126 | if method.chat_id is None or method.message_id is None: 127 | raise ValueError('Bad sent message') 128 | 129 | message = self._tg_state.get_message(int(method.chat_id), method.message_id) 130 | new_message = message.copy( 131 | update={'reply_markup': method.reply_markup}, 132 | ) 133 | self._tg_state.replace_message(new_message) 134 | return new_message 135 | 136 | METHOD_MOCKS: Mapping[Type[TelegramMethod[Any]], str] = { 137 | SendMessage: _mock_send_message.__name__, 138 | SendPhoto: _mock_send_photo.__name__, 139 | AnswerCallbackQuery: _mock_answer_callback_query.__name__, 140 | SetChatMenuButton: _mock_set_chat_menu_button.__name__, 141 | EditMessageText: _mock_edit_message_text.__name__, 142 | EditMessageReplyMarkup: _mock_edit_message_reply_markup.__name__, 143 | } 144 | 145 | async def make_request( 146 | self, 147 | bot: Bot, 148 | method: TelegramMethod[TelegramType], 149 | timeout: Optional[int] = UNSET, 150 | ) -> TelegramType: 151 | self._sent_methods.append(method) 152 | try: 153 | method_mock_attr = self.METHOD_MOCKS[type(method)] 154 | except KeyError: 155 | raise TypeError(f'Method mock for type {type(method)} is not implemented') from None 156 | 157 | return await getattr(self, method_mock_attr)(bot, method, timeout) 158 | 159 | def stream_content( 160 | self, 161 | url: str, 162 | timeout: int, 163 | chunk_size: int, 164 | raise_for_status: bool, 165 | ) -> AsyncGenerator[bytes, None]: 166 | raise NotImplementedError 167 | 168 | @property 169 | def sent_methods(self) -> Sequence[TelegramMethod[Any]]: 170 | return self._sent_methods 171 | -------------------------------------------------------------------------------- /src/aiogram_mock/tg_state.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | from collections import defaultdict 3 | from dataclasses import dataclass, replace 4 | from typing import DefaultDict, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union 5 | from uuid import uuid4 6 | 7 | from aiogram.methods import AnswerCallbackQuery 8 | from aiogram.types import ( 9 | UNSET, 10 | Chat, 11 | Document, 12 | ForceReply, 13 | InlineKeyboardMarkup, 14 | InputFile, 15 | Message, 16 | ReplyKeyboardMarkup, 17 | ) 18 | 19 | 20 | @dataclass(frozen=True) 21 | class UserState: 22 | reply_markup: Union[ReplyKeyboardMarkup, ForceReply, None] = None 23 | 24 | 25 | class TgState: 26 | def __init__(self, chats: Iterable[Chat]): 27 | self._chats = {chat.id: chat for chat in chats} 28 | self._histories: Dict[int, List[Optional[Message]]] = {chat.id: [] for chat in chats} 29 | self._message_dict: Dict[Tuple[int, int], Message] = {} 30 | self._last_update_id: int = 0 31 | self._last_callback_query_id: int = 0 32 | self._answers: Dict[str, AnswerCallbackQuery] = {} 33 | 34 | self._chat_user_state: Dict[int, UserState] = {chat.id: UserState() for chat in chats} 35 | self._selective_user_state: Dict[int, Dict[int, UserState]] = {} 36 | 37 | self._content_to_unique_id: Dict[bytes, str] = {} 38 | self._user_id_to_unique_id_to_local_id: DefaultDict[int, Dict[str, str]] = defaultdict(dict) 39 | self._user_id_to_local_id_to_unique_id: DefaultDict[int, Dict[str, str]] = defaultdict(dict) 40 | 41 | @property 42 | def chats(self) -> Mapping[int, Chat]: 43 | return self._chats 44 | 45 | def chat_history(self, chat_id: int) -> Sequence[Message]: 46 | return [message for message in self._histories[chat_id] if message is not None] 47 | 48 | def add_chat(self, chat: Chat, history: Iterable[Message] = ()) -> None: 49 | if chat.id in self._chats: 50 | raise ValueError('chat.id duplication') 51 | 52 | self._chats[chat.id] = chat 53 | self._histories[chat.id] = list(history) 54 | self._chat_user_state[chat.id] = UserState() 55 | 56 | def next_message_id(self, chat_id: int) -> int: 57 | return len(self._histories[chat_id]) 58 | 59 | def _validate_message(self, message: Message) -> None: 60 | if isinstance(message.reply_markup, InlineKeyboardMarkup): 61 | buttons = itertools.chain.from_iterable(message.reply_markup.inline_keyboard) 62 | for button in buttons: 63 | if button.callback_data is not None and len(button.callback_data) > 64: 64 | raise ValueError(f'callback_data of {button} has more than 64 chars') 65 | 66 | def add_message(self, message: Message) -> Message: 67 | if (message.chat.id, message.message_id) in self._message_dict: 68 | raise ValueError('(message.chat.id, message.message_id) duplication') 69 | 70 | self._validate_message(message) 71 | self._histories[message.chat.id].append(message) 72 | self._message_dict[(message.chat.id, message.message_id)] = message 73 | return message 74 | 75 | def get_message(self, chat_id: int, message_id: int) -> Message: 76 | return self._message_dict[(chat_id, message_id)] 77 | 78 | def replace_message(self, new_message: Message) -> None: 79 | if (new_message.chat.id, new_message.message_id) not in self._message_dict: 80 | raise KeyError('(message.chat.id, message.message_id) not exists') 81 | 82 | self._validate_message(new_message) 83 | self._histories[new_message.chat.id][new_message.message_id] = new_message 84 | self._message_dict[(new_message.chat.id, new_message.message_id)] = new_message 85 | 86 | def delete_message(self, chat_id: int, message_id: int) -> None: 87 | self._histories[chat_id][message_id] = None 88 | del self._message_dict[(chat_id, message_id)] 89 | 90 | def increment_update_id(self) -> int: 91 | self._last_update_id += 1 92 | return self._last_update_id 93 | 94 | def next_callback_query_id(self) -> str: 95 | self._last_callback_query_id += 1 96 | return str(self._last_callback_query_id) 97 | 98 | def add_answer_callback_query(self, answer: AnswerCallbackQuery) -> None: 99 | if answer.callback_query_id in self._answers: 100 | raise ValueError('callback_query_id duplication') 101 | self._answers[answer.callback_query_id] = answer 102 | 103 | def get_answer_callback_query(self, callback_query_id: str) -> AnswerCallbackQuery: 104 | return self._answers[callback_query_id] 105 | 106 | def get_user_state(self, *, chat_id: int, user_id: int) -> UserState: 107 | try: 108 | return self._selective_user_state.get(chat_id, {})[user_id] 109 | except KeyError: 110 | pass 111 | 112 | return self._chat_user_state[chat_id] 113 | 114 | def update_chat_user_state( 115 | self, 116 | chat_id: int, 117 | reply_markup: Union[ReplyKeyboardMarkup, ForceReply, None] = UNSET, 118 | ) -> None: 119 | replace_data = {} 120 | if reply_markup != UNSET: 121 | replace_data['reply_markup'] = reply_markup 122 | 123 | if chat_id in self._selective_user_state: 124 | self._selective_user_state[chat_id] = { 125 | user_id: replace(user_state, **replace_data) 126 | for user_id, user_state in self._selective_user_state[chat_id].items() 127 | } 128 | self._chat_user_state[chat_id] = replace(self._chat_user_state[chat_id], **replace_data) 129 | 130 | def update_selective_user_state( 131 | self, 132 | chat_id: int, 133 | users_ids: Iterable[int], 134 | reply_markup: Union[ReplyKeyboardMarkup, ForceReply, None] = UNSET, 135 | ) -> None: 136 | replace_data = {} 137 | if reply_markup != UNSET: 138 | replace_data['reply_markup'] = reply_markup 139 | 140 | chat_dict = self._selective_user_state[chat_id] 141 | for user_id in users_ids: 142 | chat_dict[user_id] = replace(chat_dict.get(user_id, UserState()), **replace_data) 143 | 144 | def _generate_file_unique_id(self) -> str: 145 | return str(uuid4()) 146 | 147 | def _generate_file_local_id(self, user_id: int) -> str: 148 | return f'{user_id}-{str(uuid4())}' 149 | 150 | async def _read_content(self, input_file: InputFile) -> bytes: 151 | parts = [] 152 | async for chunk in input_file: 153 | parts.append(chunk) 154 | return b''.join(parts) 155 | 156 | def _get_or_create_file_unique_id(self, content: bytes) -> str: 157 | if content in self._content_to_unique_id: 158 | return self._content_to_unique_id[content] 159 | 160 | unique_id = self._generate_file_unique_id() 161 | self._content_to_unique_id[content] = unique_id 162 | return unique_id 163 | 164 | def _get_or_create_file_local_id(self, user_id: int, unique_id: str) -> str: 165 | unique_id_to_local_id = self._user_id_to_unique_id_to_local_id[user_id] 166 | if unique_id in unique_id_to_local_id: 167 | return unique_id_to_local_id[unique_id] 168 | 169 | local_id_to_unique_id = self._user_id_to_local_id_to_unique_id[user_id] 170 | local_id = self._generate_file_local_id(user_id) 171 | unique_id_to_local_id[unique_id] = local_id 172 | local_id_to_unique_id[local_id] = unique_id 173 | return local_id 174 | 175 | async def load_file(self, user_id: int, input_file: Union[InputFile, str]) -> Document: 176 | if isinstance(input_file, str): 177 | unique_id = self._user_id_to_local_id_to_unique_id[user_id][input_file] 178 | return Document( 179 | file_id=input_file, 180 | file_unique_id=unique_id, 181 | # need to save file_name 182 | ) 183 | 184 | content = await self._read_content(input_file) 185 | unique_id = self._get_or_create_file_unique_id(content) 186 | local_id = self._get_or_create_file_local_id(user_id, unique_id) 187 | return Document( 188 | file_id=local_id, 189 | file_unique_id=unique_id, 190 | file_name=input_file.filename, 191 | ) 192 | --------------------------------------------------------------------------------