Fix typing for mocked_bot.py

This commit is contained in:
jrootjunior 2019-11-18 17:52:35 +02:00
parent ef956f7ec2
commit 7db212593b

View file

@ -1,19 +1,19 @@
from collections import deque from collections import deque
from typing import Optional from typing import TYPE_CHECKING, Deque, Optional, Type
from aiogram import Bot from aiogram import Bot
from aiogram.api.client.session.base import BaseSession from aiogram.api.client.session.base import BaseSession
from aiogram.api.methods import TelegramMethod from aiogram.api.methods import TelegramMethod
from aiogram.api.methods.base import T, Request, Response from aiogram.api.methods.base import Request, Response, T
class MockedSession(BaseSession): class MockedSession(BaseSession):
def __init__(self): def __init__(self):
super(MockedSession, self).__init__() super(MockedSession, self).__init__()
self.responses = deque() self.responses: Deque[Response[T]] = deque()
self.requests = deque() self.requests: Deque[Request] = deque()
def add_result(self, response: Response) -> Response: def add_result(self, response: Response[T]) -> Response[T]:
self.responses.append(response) self.responses.append(response)
return response return response
@ -25,18 +25,21 @@ class MockedSession(BaseSession):
async def make_request(self, token: str, method: TelegramMethod[T]) -> T: async def make_request(self, token: str, method: TelegramMethod[T]) -> T:
self.requests.append(method.build_request()) self.requests.append(method.build_request())
response: Response = self.responses.pop() response: Response[T] = self.responses.pop()
self.raise_for_status(response) self.raise_for_status(response)
return response.result return response.result # type: ignore
class MockedBot(Bot): class MockedBot(Bot):
if TYPE_CHECKING:
session: MockedSession
def __init__(self): def __init__(self):
super(MockedBot, self).__init__("TOKEN", session=MockedSession()) super(MockedBot, self).__init__("TOKEN", session=MockedSession())
def add_result_for( def add_result_for(
self, self,
method, method: Type[TelegramMethod[T]],
ok: bool, ok: bool,
result: Optional[T] = None, result: Optional[T] = None,
description: Optional[str] = None, description: Optional[str] = None,
@ -44,7 +47,7 @@ class MockedBot(Bot):
migrate_to_chat_id: Optional[int] = None, migrate_to_chat_id: Optional[int] = None,
retry_after: Optional[int] = None, retry_after: Optional[int] = None,
) -> Response[T]: ) -> Response[T]:
response = Response[method.__returning__]( response = Response[method.__returning__]( # type: ignore
ok=ok, ok=ok,
result=result, result=result,
description=description, description=description,