Repository: Zapix/mtpylon (language: python)

Issue:
Add Strategies To Handle Different Messages
MTProto has got several different message types: Encrypted/Unencrypted messages, Message containers, Service messages and Rpc Calls/Responses.  Each type has got a little bit different handling strategies. Strategy that should be used could be determined and applied in `MessageHandler`. To implement this we need:

- [ ] Describe protocol of Strategy.
- [ ] Add function that selects strategy by income message
- [ ] Apply selected strategy to function

Coz Handling Unencrypted messages has been implemented: 

- [ ] Move unencrypted message handling to strategy


Discussion / hints:

Fix patch:
diff --git a/mtpylon/message_handler.py b/mtpylon/message_handler.py
deleted file mode 100644
index a2f2592..0000000
--- a/mtpylon/message_handler.py
+++ /dev/null
@@ -1,103 +0,0 @@
-# -*- coding: utf-8 -*-
-import logging
-from typing import cast, Any, List
-from dataclasses import dataclass, field
-from functools import partial
-
-from aiohttp import web
-
-from .exceptions import InvalidMessageError, InvalidServerSalt
-from .schema import Schema
-from .serialization import CallableFunc
-from .transports import Obfuscator, TransportWrapper
-from .message_sender import MessageSender
-from .messages import MtprotoMessage, unpack_message
-from .service_schema.constructors import (
-    BadMessageNotification,
-    BadServerSalt
-)
-from .utils import get_function_name
-from .middlewares import MiddleWareFunc
-from .contextvars import income_message_var
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class MessageHandler:
-    schema: Schema
-    obfuscator: Obfuscator
-    transport_wrapper: TransportWrapper
-    message_sender: MessageSender
-    middlewares: List[MiddleWareFunc] = field(default_factory=list)
-
-    async def handle(self, request: web.Request, obfuscated_data: bytes):
-        message = await self.decrypt_message(request, obfuscated_data)
-        logger.debug(f'Received message: {message.message_id}')
-        income_message_var.set(message)
-
-        result: Any = None
-
-        try:
-            self.validate_message(message)
-        except InvalidMessageError as e:
-            msg_id = message.message_id
-            e_code = e.error_code
-            logger.error(f'Msg {msg_id}: is invalid message. code: {e_code}')
-            result = BadMessageNotification(
-                bad_msg_id=msg_id,
-                bad_msg_seqno=getattr(message, 'seqno', 0),
-                error_code=e.error_code
-            )
-        except InvalidServerSalt as e:
-            msg_id = message.message_id
-            logger.error(f'Msg {msg_id}: has bad server salt.')
-            logger.info(f'Msg {msg_id}: new server salt: {e.new_server_salt}')
-            result = BadServerSalt(
-                bad_msg_id=msg_id,
-                bad_msg_seqno=getattr(message, 'seqno', 0),
-                error_code=e.error_code,
-                new_server_salt=e.new_server_salt
-            )
-        else:
-            value = cast(CallableFunc, message.message_data)
-            func_name = get_function_name(value.func)
-            logger.info(
-                f'Msg {message.message_id}: call rpc function: {func_name}'
-            )
-
-            handler = value.func
-
-            for middleware in self.middlewares[::-1]:
-                handler = partial(middleware, handler)
-
-            result = await handler(request, **value.params)
-
-        logger.info(f'Response to message {message.message_id}')
-        await self.message_sender.send_message(request, result, response=True)
-
-    async def decrypt_message(
-            self,
-            request: web.Request,
-            obfuscated_data: bytes
-    ) -> MtprotoMessage:
-        transport_message = self.obfuscator.decrypt(obfuscated_data)
-        message_bytes = self.transport_wrapper.unwrap(transport_message)
-
-        return await unpack_message(
-            request.app['auth_key_manager'],
-            self.schema,
-            message_bytes
-        )
-
-    def validate_message(self, msg: MtprotoMessage) -> MtprotoMessage:
-        """
-        Validate message and returns it.
-
-        Raises:
-            InvalidMessageError - when wrong message has been received
-
-        :param msg:
-        :return:
-        """
-        ...
diff --git a/mtpylon/message_handler/__init__.py b/mtpylon/message_handler/__init__.py
new file mode 100644
index 0000000..33f3994
--- /dev/null
+++ b/mtpylon/message_handler/__init__.py
@@ -0,0 +1,6 @@
+# -*- coding: utf-8 -*-
+from .message_handler import MessageHandler
+
+__all__ = [
+    'MessageHandler'
+]
diff --git a/mtpylon/message_handler/message_handler.py b/mtpylon/message_handler/message_handler.py
new file mode 100644
index 0000000..fe23bb2
--- /dev/null
+++ b/mtpylon/message_handler/message_handler.py
@@ -0,0 +1,48 @@
+# -*- coding: utf-8 -*-
+import logging
+from typing import List
+from dataclasses import dataclass, field
+
+from aiohttp import web
+
+from mtpylon.schema import Schema
+from mtpylon.transports import Obfuscator, TransportWrapper
+from mtpylon.message_sender import MessageSender
+from mtpylon.messages import MtprotoMessage, unpack_message
+from mtpylon.middlewares import MiddleWareFunc
+
+from .strategies import get_handle_strategy
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MessageHandler:
+    schema: Schema
+    obfuscator: Obfuscator
+    transport_wrapper: TransportWrapper
+    message_sender: MessageSender
+    middlewares: List[MiddleWareFunc] = field(default_factory=list)
+
+    async def handle(self, request: web.Request, obfuscated_data: bytes):
+        message = await self.decrypt_message(request, obfuscated_data)
+        logger.info(f'Received message: {message.message_id}')
+        logger.debug(message)
+
+        handler = get_handle_strategy(message)
+
+        await handler(self.middlewares, self.message_sender, request, message)
+
+    async def decrypt_message(
+            self,
+            request: web.Request,
+            obfuscated_data: bytes
+    ) -> MtprotoMessage:
+        transport_message = self.obfuscator.decrypt(obfuscated_data)
+        message_bytes = self.transport_wrapper.unwrap(transport_message)
+
+        return await unpack_message(
+            request.app['auth_key_manager'],
+            self.schema,
+            message_bytes
+        )
diff --git a/mtpylon/message_handler/strategies/__init__.py b/mtpylon/message_handler/strategies/__init__.py
new file mode 100644
index 0000000..378bce6
--- /dev/null
+++ b/mtpylon/message_handler/strategies/__init__.py
@@ -0,0 +1,6 @@
+# -*- coding: utf-8 -*-
+from .selector import get_handle_strategy
+
+__all__ = [
+    'get_handle_strategy',
+]
diff --git a/mtpylon/message_handler/strategies/handle_unencrypted_message.py b/mtpylon/message_handler/strategies/handle_unencrypted_message.py
new file mode 100644
index 0000000..1f86230
--- /dev/null
+++ b/mtpylon/message_handler/strategies/handle_unencrypted_message.py
@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+import logging
+from typing import cast, List
+
+from aiohttp.web import Request
+
+from mtpylon.messages import MtprotoMessage, UnencryptedMessage
+from mtpylon.serialization import CallableFunc
+from mtpylon.middlewares import MiddleWareFunc
+from mtpylon.message_sender import MessageSender
+from mtpylon.contextvars import income_message_var
+from mtpylon.utils import get_function_name
+from mtpylon.middlewares import apply_middleware
+
+logger = logging.getLogger(__name__)
+
+
+async def handle_unencrypted_message(
+    middlewares: List[MiddleWareFunc],
+    sender: MessageSender,
+    request: Request,
+    message: MtprotoMessage,
+):
+    """
+    Handles unencrypted message. Applies all middlewares to sender function
+    and sends message
+    """
+    message = cast(UnencryptedMessage, message)
+
+    income_message_var.set(message)
+
+    value = cast(CallableFunc, message.message_data)
+    func_name = get_function_name(value.func)
+
+    logger.info(
+        f'Msg {message.message_id}: call rpc function: {func_name}'
+    )
+
+    handler = apply_middleware(middlewares, value.func)
+
+    result = await handler(request, **value.params)
+
+    logger.info(f'Response to message {message.message_id}')
+
+    await sender.send_message(request, result, response=True)
diff --git a/mtpylon/message_handler/strategies/handle_unknown_message.py b/mtpylon/message_handler/strategies/handle_unknown_message.py
new file mode 100644
index 0000000..950b9ef
--- /dev/null
+++ b/mtpylon/message_handler/strategies/handle_unknown_message.py
@@ -0,0 +1,23 @@
+# -*- coding: utf-8 -*-
+import logging
+from typing import List
+
+from aiohttp.web import Request
+
+from mtpylon.messages import MtprotoMessage
+from mtpylon.middlewares import MiddleWareFunc
+from mtpylon.message_sender import MessageSender
+
+logger = logging.getLogger(__name__)
+
+
+async def handle_unknown_message(
+    middlewares: List[MiddleWareFunc],
+    sender: MessageSender,
+    request: Request,
+    message: MtprotoMessage,
+):
+    """
+    Logs that we don't know how to handle this message
+    """
+    logger.warning(f'Unknown message: {message}')
diff --git a/mtpylon/message_handler/strategies/selector.py b/mtpylon/message_handler/strategies/selector.py
new file mode 100644
index 0000000..c7464d0
--- /dev/null
+++ b/mtpylon/message_handler/strategies/selector.py
@@ -0,0 +1,16 @@
+# -*- coding: utf-8 -*-
+from mtpylon.messages import MtprotoMessage, UnencryptedMessage
+
+from .types import HandleStrategy
+from .handle_unknown_message import handle_unknown_message
+from .handle_unencrypted_message import handle_unencrypted_message
+
+
+def get_handle_strategy(message: MtprotoMessage) -> HandleStrategy:
+    """
+    Selects how message will be handled by message
+    """
+    if isinstance(message, UnencryptedMessage):
+        return handle_unencrypted_message
+
+    return handle_unknown_message
diff --git a/mtpylon/message_handler/strategies/types.py b/mtpylon/message_handler/strategies/types.py
new file mode 100644
index 0000000..078028a
--- /dev/null
+++ b/mtpylon/message_handler/strategies/types.py
@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+from typing import Callable, Awaitable, List
+
+from mypy_extensions import Arg
+from aiohttp.web import Request
+
+
+from mtpylon.messages import MtprotoMessage
+from mtpylon.middlewares import MiddleWareFunc
+from mtpylon.message_sender import MessageSender
+
+
+HandleStrategy = Callable[
+    [
+        Arg(List[MiddleWareFunc], 'middlewares'),  # noqa: F821
+        Arg(MessageSender, 'sender'),  # noqa: F821
+        Arg(Request, 'request'),  # noqa: F821
+        Arg(MtprotoMessage, 'message'),  # noqa: F821
+    ],
+    Awaitable[None]
+]
diff --git a/mtpylon/middlewares/__init__.py b/mtpylon/middlewares/__init__.py
index 2eb7bce..a3d1f98 100644
--- a/mtpylon/middlewares/__init__.py
+++ b/mtpylon/middlewares/__init__.py
@@ -2,6 +2,7 @@
 from .types import MiddleWareFunc
 from .set_server_salt import set_server_salt
 from .set_session_id import set_session_id
+from .apply import apply_middleware
 
 BASIC_MIDDLEWARES = [
     set_session_id,
@@ -10,5 +11,6 @@
 
 __all__ = [
     'MiddleWareFunc',
-    'BASIC_MIDDLEWARES'
+    'BASIC_MIDDLEWARES',
+    'apply_middleware',
 ]
diff --git a/mtpylon/middlewares/apply.py b/mtpylon/middlewares/apply.py
new file mode 100644
index 0000000..0dd39e5
--- /dev/null
+++ b/mtpylon/middlewares/apply.py
@@ -0,0 +1,15 @@
+# -*- coding: utf-8 -*-
+from typing import List
+from functools import partial
+
+from .types import MiddleWareFunc, Handler
+
+
+def apply_middleware(
+    middlewares: List[MiddleWareFunc],
+    handler: Handler
+) -> Handler:
+    for middleware in middlewares[::-1]:
+        handler = partial(middleware, handler)
+
+    return handler


Test patch:
diff --git a/tests/test_message_handler/__init__.py b/tests/test_message_handler/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/tests/test_message_handler/__init__.py
@@ -0,0 +1,1 @@
+# -*- coding: utf-8 -*-
diff --git a/tests/test_message_handler.py b/tests/test_message_handler/test_message_handler.py
similarity index 50%
rename from tests/test_message_handler.py
rename to tests/test_message_handler/test_message_handler.py
index cc2fc61..29aa6b4 100644
--- a/tests/test_message_handler.py
+++ b/tests/test_message_handler/test_message_handler.py
@@ -3,21 +3,16 @@
 
 import pytest
 
-from tests.simpleschema import schema, set_task, Task
-
 from mtpylon import long
 from mtpylon.serialization import CallableFunc
 from mtpylon.message_handler import MessageHandler
-from mtpylon.exceptions import InvalidMessageError, InvalidServerSalt
 from mtpylon.messages import UnencryptedMessage, Message
-from mtpylon.service_schema.constructors import (
-    BadMessageNotification,
-    BadServerSalt
-)
 from mtpylon.contextvars import (
     income_message_var
 )
 
+from tests.simpleschema import schema, set_task, Task
+
 
 @pytest.mark.asyncio
 async def test_unpack_unencyprted_message_correct():
@@ -38,12 +33,14 @@ async def test_unpack_unencyprted_message_correct():
             func=set_task,
             params={'content': 'hello world'}
         )
-
     )
 
     unpack_message = AsyncMock(return_value=message)
 
-    with patch('mtpylon.message_handler.unpack_message', unpack_message):
+    with patch(
+        'mtpylon.message_handler.message_handler.unpack_message',
+        unpack_message
+    ):
         handler = MessageHandler(
             schema=schema,
             obfuscator=obfuscator,
@@ -93,7 +90,10 @@ async def test_unpack_encrypted_message_correct():
 
     unpack_message = AsyncMock(return_value=message)
 
-    with patch('mtpylon.message_handler.unpack_message', unpack_message):
+    with patch(
+        'mtpylon.message_handler.message_handler.unpack_message',
+        unpack_message
+    ):
         handler = MessageHandler(
             schema=schema,
             obfuscator=obfuscator,
@@ -107,115 +107,6 @@ async def test_unpack_encrypted_message_correct():
         assert transport_wrapper.unwrap.called
         unpack_message.assert_awaited()
 
-        message_sender.send_message.assert_awaited()
-        task = message_sender.send_message.await_args[0][1]
-        assert isinstance(task, Task)
-        assert task.id == 1
-        assert task.content == 'hello world'
-        assert income_message_var.get() == message
-
-
-@pytest.mark.asyncio
-async def test_invalid_message_id():
-    obfuscator = MagicMock()
-    obfuscator.decrypt__return_value = b'decrypted data'
-    transport_wrapper = MagicMock()
-    transport_wrapper.unwrap__return_value = b'unwrapped data'
-    message_sender = MagicMock(
-        send_message=AsyncMock()
-    )
-    request = MagicMock()
-
-    msg_id = long(0x51e57ac42770964a)
-
-    unpack_message = AsyncMock(
-        return_value=UnencryptedMessage(
-            message_id=msg_id,
-            message_data=CallableFunc(
-                func=set_task,
-                params={'content': 'hello world'}
-            )
-        )
-    )
-
-    validate_message = MagicMock(side_effect=InvalidMessageError(
-        error_code=64
-    ))
-
-    with patch('mtpylon.message_handler.unpack_message', unpack_message):
-        handler = MessageHandler(
-            schema=schema,
-            obfuscator=obfuscator,
-            transport_wrapper=transport_wrapper,
-            message_sender=message_sender,
-        )
-
-        with patch.object(handler, 'validate_message', validate_message):
-            await handler.handle(request, b'obfuscated message')
-
-            assert obfuscator.decrypt.called
-            assert transport_wrapper.unwrap.called
-            unpack_message.assert_awaited()
-
-            message_sender.send_message.assert_awaited()
-            error = message_sender.send_message.await_args[0][1]
-            assert isinstance(error, BadMessageNotification)
-            assert error.bad_msg_id == msg_id
-            assert error.bad_msg_seqno == 0
-            assert error.error_code == 64
-
-
-@pytest.mark.asyncio
-async def test_invalid_server_salt():
-    obfuscator = MagicMock()
-    obfuscator.decrypt__return_value = b'decrypted data'
-    transport_wrapper = MagicMock()
-    transport_wrapper.unwrap__return_value = b'unwrapped data'
-    message_sender = MagicMock(
-        send_message=AsyncMock()
-    )
-    request = MagicMock()
-
-    msg_id = long(0x51e57ac42770964a)
-
-    unpack_message = AsyncMock(
-        return_value=UnencryptedMessage(
-            message_id=msg_id,
-            message_data=CallableFunc(
-                func=set_task,
-                params={'content': 'hello world'}
-            )
-        )
-    )
-
-    validate_message = MagicMock(side_effect=InvalidServerSalt(
-        error_code=48,
-        new_server_salt=long(0xacab1312)
-    ))
-
-    with patch('mtpylon.message_handler.unpack_message', unpack_message):
-        handler = MessageHandler(
-            schema=schema,
-            obfuscator=obfuscator,
-            transport_wrapper=transport_wrapper,
-            message_sender=message_sender
-        )
-
-        with patch.object(handler, 'validate_message', validate_message):
-            await handler.handle(request, b'obfuscated message')
-
-            assert obfuscator.decrypt.called
-            assert transport_wrapper.unwrap.called
-            unpack_message.assert_awaited()
-
-            message_sender.send_message.assert_awaited()
-            error = message_sender.send_message.await_args[0][1]
-            assert isinstance(error, BadServerSalt)
-            assert error.bad_msg_id == msg_id
-            assert error.bad_msg_seqno == 0
-            assert error.error_code == 48
-            assert error.new_server_salt == long(0xacab1312)
-
 
 @pytest.mark.asyncio
 async def test_pass_middlewares():
@@ -247,7 +138,10 @@ async def simple_middleware(handler, request, **params):
         )
     )
 
-    with patch('mtpylon.message_handler.unpack_message', unpack_message):
+    with patch(
+        'mtpylon.message_handler.message_handler.unpack_message',
+        unpack_message
+    ):
         handler = MessageHandler(
             schema=schema,
             obfuscator=obfuscator,
diff --git a/tests/test_message_handler/test_strategies/__init__.py b/tests/test_message_handler/test_strategies/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/tests/test_message_handler/test_strategies/__init__.py
@@ -0,0 +1,1 @@
+# -*- coding: utf-8 -*-
diff --git a/tests/test_message_handler/test_strategies/test_handle_unencrypted_message.py b/tests/test_message_handler/test_strategies/test_handle_unencrypted_message.py
new file mode 100644
index 0000000..39ebf39
--- /dev/null
+++ b/tests/test_message_handler/test_strategies/test_handle_unencrypted_message.py
@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from mtpylon import long
+from mtpylon.serialization import CallableFunc
+from mtpylon.messages import UnencryptedMessage
+from mtpylon.contextvars import income_message_var
+from mtpylon.message_handler.strategies.handle_unencrypted_message import (
+    handle_unencrypted_message
+)
+
+
+from tests.simpleschema import set_task, Task
+
+
+@pytest.mark.asyncio
+async def test_handle_unencrypted_message():
+    msg_id = long(0x51e57ac42770964a)
+    message = UnencryptedMessage(
+        message_id=msg_id,
+        message_data=CallableFunc(
+            func=set_task,
+            params={'content': 'hello world!'}
+        ),
+    )
+
+    middlewares = []
+
+    request = MagicMock()
+
+    sender = MagicMock(send_message=AsyncMock())
+
+    await handle_unencrypted_message(
+        middlewares,
+        sender,
+        request,
+        message
+    )
+
+    sender.send_message.assert_awaited()
+
+    task = sender.send_message.await_args[0][1]
+    assert isinstance(task, Task)
+    assert task.id == 1
+    assert task.content == 'hello world!'
+
+    assert income_message_var.get() == message
diff --git a/tests/test_message_handler/test_strategies/test_handle_unknown_message.py b/tests/test_message_handler/test_strategies/test_handle_unknown_message.py
new file mode 100644
index 0000000..f26b076
--- /dev/null
+++ b/tests/test_message_handler/test_strategies/test_handle_unknown_message.py
@@ -0,0 +1,36 @@
+# -*- coding: utf-8 -*-
+from unittest.mock import MagicMock, AsyncMock, patch
+
+import pytest
+
+from mtpylon import long
+from mtpylon.messages import Message
+from mtpylon.message_handler.strategies.handle_unknown_message import (
+    handle_unknown_message
+)
+
+
+@pytest.mark.asyncio
+async def test_handle_unknown_message():
+    msg_id = long(0x51e57ac42770964a)
+
+    message = Message(
+        message_id=msg_id,
+        session_id=long(1),
+        salt=long(2),
+        seq_no=0,
+        message_data='Wrong message data'
+    )
+
+    request = MagicMock()
+    sender = MagicMock(send_message=AsyncMock())
+    logger = MagicMock()
+
+    with patch(
+        'mtpylon.message_handler.strategies.handle_unknown_message.logger',
+        logger,
+    ):
+        await handle_unknown_message([], sender, request, message)
+
+    logger.warning.assert_called()
+    sender.send_message.assert_not_called()
diff --git a/tests/test_message_handler/test_strategies/test_selector.py b/tests/test_message_handler/test_strategies/test_selector.py
new file mode 100644
index 0000000..aa3c7c2
--- /dev/null
+++ b/tests/test_message_handler/test_strategies/test_selector.py
@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+from mtpylon import long
+from mtpylon.serialization import CallableFunc
+from mtpylon.messages import Message, UnencryptedMessage
+from mtpylon.message_handler.strategies import get_handle_strategy
+from mtpylon.message_handler.strategies.handle_unknown_message import (
+    handle_unknown_message
+)
+from mtpylon.message_handler.strategies.handle_unencrypted_message import (
+    handle_unencrypted_message
+)
+
+from tests.simpleschema import set_task
+
+
+def test_get_unknonwn_handler():
+    msg_id = long(0x51e57ac42770964a)
+    message = Message(
+        message_id=msg_id,
+        session_id=long(1),
+        salt=long(2),
+        seq_no=0,
+        message_data='Wrong message data'
+    )
+
+    handler = get_handle_strategy(message)
+
+    assert handler == handle_unknown_message
+
+
+def test_get_unencrypted_handler():
+    msg_id = long(0x51e57ac42770964a)
+    message = UnencryptedMessage(
+        message_id=msg_id,
+        message_data=CallableFunc(
+            func=set_task,
+            params={'content': 'hello world!'}
+        ),
+    )
+
+    handler = get_handle_strategy(message)
+
+    assert handler == handle_unencrypted_message
diff --git a/tests/test_middlewares/test_apply.py b/tests/test_middlewares/test_apply.py
new file mode 100644
index 0000000..416a5fb
--- /dev/null
+++ b/tests/test_middlewares/test_apply.py
@@ -0,0 +1,30 @@
+# -*- coding: utf-8 -*-
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from mtpylon.middlewares import apply_middleware
+
+from tests.simpleschema import set_task, Task
+
+
+@pytest.mark.asyncio
+async def test_apply_middleware():
+
+    async def simple_middleware(handler, request, **params):
+        return await handler(request, **params)
+
+    middleware1 = AsyncMock(side_effect=simple_middleware)
+    middleware2 = AsyncMock(side_effect=simple_middleware)
+
+    handler = apply_middleware([middleware1, middleware2], set_task)
+
+    request = MagicMock()
+
+    result = await handler(request, content='hello world')
+
+    assert isinstance(result, Task)
+    assert result.content == 'hello world'
+
+    middleware1.assert_awaited()
+    middleware2.assert_awaited()