aiogram/tests/mocked_bot.py
Alex Root Junior ec7da0f678
Migrate from Black to Ruff (#1750)
* Migrate from Black to Ruff and reformat code with enabling additional linter checks

* Add changelog for migration to Ruff as formatter and linter

* Add type ignores for specific attributes and replace tuple with set for chat type check

* Remove file from another changes
2026-01-04 21:34:08 +02:00

96 lines
2.9 KiB
Python

from collections import deque
from collections.abc import AsyncGenerator
from typing import TYPE_CHECKING, Any
from aiogram import Bot
from aiogram.client.session.base import BaseSession
from aiogram.methods import TelegramMethod
from aiogram.methods.base import Response, TelegramType
from aiogram.types import UNSET_PARSE_MODE, ResponseParameters, User
class MockedSession(BaseSession):
def __init__(self):
super().__init__()
self.responses: deque[Response[TelegramType]] = deque()
self.requests: deque[TelegramMethod[TelegramType]] = deque()
self.closed = True
def add_result(self, response: Response[TelegramType]) -> Response[TelegramType]:
self.responses.append(response)
return response
def get_request(self) -> TelegramMethod[TelegramType]:
return self.requests.pop()
async def close(self):
self.closed = True
async def make_request(
self,
bot: Bot,
method: TelegramMethod[TelegramType],
timeout: int | None = UNSET_PARSE_MODE,
) -> TelegramType:
self.closed = False
self.requests.append(method)
response: Response[TelegramType] = self.responses.pop()
self.check_response(
bot=bot,
method=method,
status_code=response.error_code,
content=response.model_dump_json(),
)
return response.result # type: ignore
async def stream_content(
self,
url: str,
headers: dict[str, Any] | None = None,
timeout: int = 30,
chunk_size: int = 65536,
raise_for_status: bool = True,
) -> AsyncGenerator[bytes, None]: # pragma: no cover
yield b""
class MockedBot(Bot):
if TYPE_CHECKING:
session: MockedSession
def __init__(self, **kwargs):
super().__init__(kwargs.pop("token", "42:TEST"), session=MockedSession(), **kwargs)
self._me = User(
id=self.id,
is_bot=True,
first_name="FirstName",
last_name="LastName",
username="tbot",
language_code="uk-UA",
)
def add_result_for(
self,
method: type[TelegramMethod[TelegramType]],
ok: bool,
result: TelegramType = None,
description: str | None = None,
error_code: int = 200,
migrate_to_chat_id: int | None = None,
retry_after: int | None = None,
) -> Response[TelegramType]:
response = Response[method.__returning__]( # type: ignore
ok=ok,
result=result,
description=description,
error_code=error_code,
parameters=ResponseParameters(
migrate_to_chat_id=migrate_to_chat_id,
retry_after=retry_after,
),
)
self.session.add_result(response)
return response
def get_request(self) -> TelegramMethod[TelegramType]:
return self.session.get_request()