From 37d46d92500284127fc5185e29003b61f0c7d751 Mon Sep 17 00:00:00 2001
From: Oleg A <t0rr@mail.ru>
Date: Thu, 15 Nov 2018 10:29:59 +0300
Subject: [PATCH] test errors_handler args

---
 aiogram/dispatcher/dispatcher.py | 20 +++++--
 tests/test_dispatcher.py         | 99 ++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_dispatcher.py

diff --git a/aiogram/dispatcher/dispatcher.py b/aiogram/dispatcher/dispatcher.py
index 3f7c0d9b..db9e007e 100644
--- a/aiogram/dispatcher/dispatcher.py
+++ b/aiogram/dispatcher/dispatcher.py
@@ -4,7 +4,7 @@ import itertools
 import logging
 import time
 import typing
-from inspect import signature
+import inspect
 
 from .filters import Command, ContentTypeFilter, ExceptionsFilter, FiltersFactory, FuncFilter, HashTag, Regexp, \
     RegexpCommandsFilter, StateFilter, Text
@@ -814,10 +814,20 @@ class Dispatcher(DataMixin, ContextInstanceMixin):
         :param run_task: run callback in task (no wait results)
         """
 
-        # Check number of arguments of the callback, cause only two arguments accepted.
-        sig = signature(callback)
-        if len(sig.parameters) != 2:
-            raise RuntimeError('Errors handlers should accept only two arguments (current update and exception)')
+        # Check number of arguments of the callback, cause only two arguments accepted (except self, *args and **kwargs)
+        signature = inspect.signature(callback)
+        arguments = inspect.getfullargspec(callback)
+        amount_of_params = 2
+
+        if arguments.varargs:
+            amount_of_params += 1
+
+        if arguments.varkw:
+            amount_of_params += 1
+
+        if len(signature.parameters) != amount_of_params:
+            raise RuntimeError(f'Errors handlers should accept only 2 arguments '
+                               f'(current update and exception). Also you may use *args and **kwargs after that.')
 
         filters_set = self.filters_factory.resolve(self.errors_handlers,
                                                    *custom_filters,
diff --git a/tests/test_dispatcher.py b/tests/test_dispatcher.py
new file mode 100644
index 00000000..27ede047
--- /dev/null
+++ b/tests/test_dispatcher.py
@@ -0,0 +1,99 @@
+from asyncio import BaseEventLoop
+
+import pytest
+
+from aiogram import Bot, types, Dispatcher
+from aiogram.dispatcher.handler import Handler
+
+pytestmark = pytest.mark.asyncio
+TOKEN = '123456789:AABBCCDDEEFFaabbccddeeff-1234567890'
+
+
+@pytest.yield_fixture()
+async def bot(event_loop: BaseEventLoop):
+    """ Bot fixture """
+    _bot = Bot(TOKEN, loop=event_loop, parse_mode=types.ParseMode.HTML)
+    yield _bot
+    await _bot.close()
+
+
+@pytest.yield_fixture()
+async def dp(event_loop: BaseEventLoop, bot: Bot):
+    """
+    Dispatcher fixture
+    :rtype: Dispatcher
+    """
+    _dp = Dispatcher(bot=bot, loop=event_loop)
+    yield _dp
+
+
+class TestErrorsHandlerArgs:
+    @staticmethod
+    def get_handler_from(handler_category: Handler):
+        filters, handler = handler_category.handlers.pop()
+        return handler
+
+    def method_with_3_args(self, two, three):
+        pass
+
+    def method_with_4_args(self, two, three, four):
+        pass
+
+    def method_with_varargs(self, two, three, *args):
+        pass
+
+    def method_with_varkw(self, two, three, **kwargs):
+        pass
+
+    def method_with_varargs_and_varkw(self, two, three, *args, **kwargs):
+        pass
+
+    async def test_errors_handler_func_with_2_args(self, dp: Dispatcher):
+        """ Test handler as func with 2 args """
+
+        async def func(one, two):
+            pass
+
+        dp.register_errors_handler(func)
+        assert func == self.get_handler_from(dp.errors_handlers)
+
+    async def test_errors_handler_func_with_3_args(self, dp: Dispatcher):
+        """ Test handler with 3 args """
+
+        async def func(one, two, three):
+            pass
+
+        with pytest.raises(RuntimeError):
+            dp.register_errors_handler(func)
+
+        assert len(dp.errors_handlers.handlers) == 0
+
+    async def test_errors_handler_method_with_3_args(self, dp: Dispatcher):
+        """ Test handler as class method with 3 args """
+
+        dp.register_errors_handler(self.method_with_3_args)
+        assert self.method_with_3_args == self.get_handler_from(dp.errors_handlers)
+
+    async def test_errors_handler_method_with_4_args(self, dp: Dispatcher):
+        """ Test handler as class method with 4 args """
+
+        with pytest.raises(RuntimeError):
+            dp.register_errors_handler(self.method_with_4_args)
+
+    async def test_errors_handler_method_with_varargs(self, dp: Dispatcher):
+        """ Test handler as class method with 2 args (except self) and *args """
+
+        dp.register_errors_handler(self.method_with_varargs)
+        assert self.method_with_varargs == self.get_handler_from(dp.errors_handlers)
+
+    async def test_errors_handler_method_with_varkw(self, dp: Dispatcher):
+        """ Test handler as class method with 2 args (except self) and **kwargs """
+
+        dp.register_errors_handler(self.method_with_varkw)
+        assert self.method_with_varkw == self.get_handler_from(dp.errors_handlers)
+
+    async def test_errors_handler_method_with_varargs_and_varkw(self, dp: Dispatcher):
+        """ Test handler as class method with 2 args (except self), *args and **kwargs """
+
+        dp.register_errors_handler(self.method_with_varargs_and_varkw)
+        assert self.method_with_varargs_and_varkw == self.get_handler_from(dp.errors_handlers)