From d247622df6ed32395285e86340acef9a41a70caa Mon Sep 17 00:00:00 2001 From: darksidecat Date: Mon, 4 Oct 2021 14:54:12 +0300 Subject: [PATCH] add middleware for logging outgoing requests --- aiogram/client/session/base.py | 13 ++++-- .../client/session/middlewares/__init__.py | 0 aiogram/client/session/middlewares/base.py | 37 ++++++++++++++++ .../session/middlewares/request_logging.py | 39 +++++++++++++++++ .../test_session/test_middlewares/__init__.py | 0 .../test_middlewares/test_request_logging.py | 42 +++++++++++++++++++ 6 files changed, 127 insertions(+), 4 deletions(-) create mode 100644 aiogram/client/session/middlewares/__init__.py create mode 100644 aiogram/client/session/middlewares/base.py create mode 100644 aiogram/client/session/middlewares/request_logging.py create mode 100644 tests/test_api/test_client/test_session/test_middlewares/__init__.py create mode 100644 tests/test_api/test_client/test_session/test_middlewares/test_request_logging.py diff --git a/aiogram/client/session/base.py b/aiogram/client/session/base.py index 3a845cfd..e161e7ba 100644 --- a/aiogram/client/session/base.py +++ b/aiogram/client/session/base.py @@ -39,6 +39,7 @@ from ...methods import Response, TelegramMethod from ...methods.base import TelegramType from ...types import UNSET, TelegramObject from ..telegram import PRODUCTION, TelegramAPIServer +from .middlewares.base import BaseRequestMiddleware if TYPE_CHECKING: from ..bot import Bot @@ -48,15 +49,19 @@ _JsonDumps = Callable[..., str] NextRequestMiddlewareType = Callable[ ["Bot", TelegramMethod[TelegramObject]], Awaitable[Response[TelegramObject]] ] -RequestMiddlewareType = Callable[ - [NextRequestMiddlewareType, "Bot", TelegramMethod[TelegramType]], - Awaitable[Response[TelegramType]], + +RequestMiddlewareType = Union[ + BaseRequestMiddleware, + Callable[ + [NextRequestMiddlewareType, "Bot", TelegramMethod[TelegramType]], + Awaitable[Response[TelegramType]], + ], ] class BaseSession(abc.ABC): api: Default[TelegramAPIServer] = Default(PRODUCTION) - """Telegra Bot API URL patterns""" + """Telegram Bot API URL patterns""" json_loads: Default[_JsonLoads] = Default(json.loads) """JSON loader""" json_dumps: Default[_JsonDumps] = Default(json.dumps) diff --git a/aiogram/client/session/middlewares/__init__.py b/aiogram/client/session/middlewares/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/aiogram/client/session/middlewares/base.py b/aiogram/client/session/middlewares/base.py new file mode 100644 index 00000000..54296d84 --- /dev/null +++ b/aiogram/client/session/middlewares/base.py @@ -0,0 +1,37 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Awaitable, Callable + +from aiogram.methods import Response, TelegramMethod +from aiogram.types import TelegramObject + +if TYPE_CHECKING: + from ...bot import Bot + + +NextRequestMiddlewareType = Callable[ + ["Bot", TelegramMethod[TelegramObject]], Awaitable[Response[TelegramObject]] +] + + +class BaseRequestMiddleware(ABC): + """ + Generic middleware class + """ + + @abstractmethod + async def __call__( + self, + bot: "Bot", + method: TelegramMethod[TelegramObject], + make_request: NextRequestMiddlewareType, + ) -> Response[TelegramObject]: + """ + Execute middleware + + :param bot: bot for request making + :param method: Request method (Subclass of :class:`aiogram.methods.base.TelegramMethod`) + :param make_request: Wrapped make_request in middlewares chain + + :return: :class:`aiogram.methods.Response` + """ + pass diff --git a/aiogram/client/session/middlewares/request_logging.py b/aiogram/client/session/middlewares/request_logging.py new file mode 100644 index 00000000..2d77bf5b --- /dev/null +++ b/aiogram/client/session/middlewares/request_logging.py @@ -0,0 +1,39 @@ +import logging +from typing import TYPE_CHECKING, Awaitable, Callable, List, Optional, Type + +from aiogram import loggers +from aiogram.methods import TelegramMethod +from aiogram.methods.base import Response +from aiogram.types import TelegramObject + +from .base import BaseRequestMiddleware + +if TYPE_CHECKING: + from ...bot import Bot + +NextRequestMiddlewareType = Callable[ + ["Bot", TelegramMethod[TelegramObject]], Awaitable[Response[TelegramObject]] +] + +logger = logging.getLogger(__name__) + + +class RequestLogging(BaseRequestMiddleware): + def __init__( + self, ignore_methods: Optional[List[Type[TelegramMethod[TelegramObject]]]] = None + ): + self.ignore_methods = ignore_methods if ignore_methods else [] + + async def __call__( + self, + bot: "Bot", + method: TelegramMethod[TelegramObject], + make_request: NextRequestMiddlewareType, + ) -> Response[TelegramObject]: + if type(method) not in self.ignore_methods: + loggers.middlewares.info( + "Make request with method=%s by bot id=%d", + method.__class__.__name__, + bot.id, + ) + return await make_request(bot, method) diff --git a/tests/test_api/test_client/test_session/test_middlewares/__init__.py b/tests/test_api/test_client/test_session/test_middlewares/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_api/test_client/test_session/test_middlewares/test_request_logging.py b/tests/test_api/test_client/test_session/test_middlewares/test_request_logging.py new file mode 100644 index 00000000..c1d6b5b3 --- /dev/null +++ b/tests/test_api/test_client/test_session/test_middlewares/test_request_logging.py @@ -0,0 +1,42 @@ +import datetime +import logging + +import pytest + +from aiogram.client.session.middlewares.request_logging import RequestLogging +from aiogram.methods import GetMe, SendMessage +from aiogram.types import Chat, Message, User +from tests.mocked_bot import MockedBot + +pytestmark = pytest.mark.asyncio + + +class TestRequestLogging: + async def test_use_middleware(self, bot: MockedBot, caplog): + caplog.set_level(logging.INFO) + bot.session.middleware(RequestLogging()) + + bot.add_result_for(GetMe, ok=True, result=User(id=42, is_bot=True, first_name="Test")) + assert await bot.get_me() + assert "Make request with method=GetMe by bot id=42" in caplog.text + + async def test_ignore_methods(self, bot: MockedBot, caplog): + caplog.set_level(logging.INFO) + bot.session.middleware(RequestLogging(ignore_methods=[GetMe])) + + bot.add_result_for(GetMe, ok=True, result=User(id=42, is_bot=True, first_name="Test")) + assert await bot.get_me() + assert "Make request with method=GetMe by bot id=42" not in caplog.text + + bot.add_result_for( + SendMessage, + ok=True, + result=Message( + message_id=42, + date=datetime.datetime.now(), + text="test", + chat=Chat(id=42, type="private"), + ), + ) + assert await bot.send_message(chat_id=1, text="Test") + assert "Make request with method=SendMessage by bot id=42" in caplog.text