Repository: Zapix/mtpylon (language: python)

Issue:
Sesssion Manager
- [x] Describe `SessionStorageProtocol`, `SessionSubjectProtocol` , `SessionObserverProtocol`
- [x] Store info about created sessions(session ids)
- [x] Add Subject that notifies when session created, destroyed
- [x] Checks that subject has been added to app


Discussion / hints:

Fix patch:
diff --git a/mtpylon/aiohandlers/websockets.py b/mtpylon/aiohandlers/websockets.py
index fd33034..b8372c9 100644
--- a/mtpylon/aiohandlers/websockets.py
+++ b/mtpylon/aiohandlers/websockets.py
@@ -2,9 +2,10 @@
 import logging
 from functools import partial
 from typing import cast, Optional
+from dataclasses import dataclass
 
 from aiohttp import WSMsgType
-from aiohttp.web import Request, WebSocketResponse
+from aiohttp.web import Request, WebSocketResponse, Application
 
 from mtpylon.aiohandlers.types import WebSocketHandler
 from mtpylon.schema import Schema
@@ -18,6 +19,48 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True)
+class SharedResourceCheck:
+    resource_name: str
+    error: str
+
+
+SHARED_RESOURCE_CHECK_LIST = [
+    SharedResourceCheck(
+        resource_name='rsa_manager',
+        error='Rsa manager should be set',
+    ),
+    SharedResourceCheck(
+        resource_name='auth_key_manager',
+        error='Auth key manager should be set',
+    ),
+    SharedResourceCheck(
+        resource_name='dh_prime_generator',
+        error='DH prime generator should be set',
+    ),
+    SharedResourceCheck(
+        resource_name='server_salt_manager',
+        error='Server salt manager should be set'
+    ),
+    SharedResourceCheck(
+        resource_name='session_subject',
+        error='Session subject should be set',
+    )
+]
+
+
+def validate_shared_resources(app: Application):
+    """
+    Validates that all shared resources properly configures
+
+    Raises:
+        ValueError - if some of resources hasn't been configured
+    """
+    for check in SHARED_RESOURCE_CHECK_LIST:
+        if check.resource_name not in app:
+            raise ValueError(check.error)
+
+
 async def ws_handler(request: Request, schema: Schema) -> WebSocketResponse:
     """
     Websocket handler that works with MTProto protocol.
@@ -32,23 +75,10 @@ async def ws_handler(request: Request, schema: Schema) -> WebSocketResponse:
     ws = WebSocketResponse()
     await ws.prepare(request)
 
-    if 'rsa_manager' not in request.app:
-        logger.error('Rsa manager should be set')
-        await ws.close()
-        return ws
-
-    if 'auth_key_manager' not in request.app:
-        logger.error('Auth key manager should be set')
-        await ws.close()
-        return ws
-
-    if 'dh_prime_generator' not in request.app:
-        logger.error('DH prime generator should be set')
-        await ws.close()
-        return ws
-
-    if 'server_salt_manager' not in request.app:
-        logger.error('Server salt manager should be set')
+    try:
+        validate_shared_resources(request.app)
+    except ValueError as e:
+        logger.error(e)
         await ws.close()
         return ws
 
diff --git a/mtpylon/sessions/__init__.py b/mtpylon/sessions/__init__.py
new file mode 100644
index 0000000..601ae51
--- /dev/null
+++ b/mtpylon/sessions/__init__.py
@@ -0,0 +1,10 @@
+# -*- coding: utf-8 -*-
+from .in_memory_session_storage import InMemorySessionStorage
+from .session_subject import SessionSubject
+from .session_event import SessionEvent
+
+__all__ = [
+    'InMemorySessionStorage',
+    'SessionSubject',
+    'SessionEvent',
+]
diff --git a/mtpylon/sessions/in_memory_session_storage.py b/mtpylon/sessions/in_memory_session_storage.py
new file mode 100644
index 0000000..c1de05d
--- /dev/null
+++ b/mtpylon/sessions/in_memory_session_storage.py
@@ -0,0 +1,28 @@
+# -*- coding: utf-8 -*-
+from typing import Set
+from asyncio import Lock
+
+from mtpylon.types import long
+from .session_storage_protocol import SessionStorageProtocol
+
+
+class InMemorySessionStorage(SessionStorageProtocol):
+    """
+    Stores available session in session set
+    """
+
+    def __init__(self):
+        self._storage: Set[long] = set()
+        self._lock = Lock()
+
+    async def create_session(self, session_id: long):
+        async with self._lock:
+            self._storage.add(session_id)
+
+    async def has_session(self, session_id: long) -> bool:
+        return session_id in self._storage
+
+    async def destroy_session(self, session_id: long):
+        if session_id in self._storage:
+            async with self._lock:
+                self._storage.remove(session_id)
diff --git a/mtpylon/sessions/session_event.py b/mtpylon/sessions/session_event.py
new file mode 100644
index 0000000..3f63154
--- /dev/null
+++ b/mtpylon/sessions/session_event.py
@@ -0,0 +1,14 @@
+# -*- coding: utf-8 -*-
+from typing import Literal
+from dataclasses import dataclass
+
+from mtpylon.types import long
+
+
+EventType = Literal['created', 'destroyed']
+
+
+@dataclass(frozen=True)
+class SessionEvent:
+    type: EventType
+    session_id: long
diff --git a/mtpylon/sessions/session_oberver_protocol.py b/mtpylon/sessions/session_oberver_protocol.py
new file mode 100644
index 0000000..03124c0
--- /dev/null
+++ b/mtpylon/sessions/session_oberver_protocol.py
@@ -0,0 +1,8 @@
+# -*- coding: utf-8 -*-
+from .session_event import SessionEvent
+
+
+class SessionObserverProtocol:
+
+    async def update(self, event: SessionEvent):  # pragma: nocover
+        ...
diff --git a/mtpylon/sessions/session_storage_protocol.py b/mtpylon/sessions/session_storage_protocol.py
new file mode 100644
index 0000000..6de98fc
--- /dev/null
+++ b/mtpylon/sessions/session_storage_protocol.py
@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+from typing import Protocol
+
+from mtpylon import long
+
+
+class SessionStorageProtocol(Protocol):
+
+    async def create_session(self, session_id: long):  # pragma: nocover
+        """
+        Saves info about new session(session_id)
+        """
+        ...
+
+    async def has_session(self, session_id: long) -> bool:  # pragma: nocover
+        """
+        Checks has session_id been stored or not
+        """
+        ...
+
+    async def destroy_session(self, session_id: long):  # pragma: nocover
+        """
+        Removes info about stored
+        """
+        ...
diff --git a/mtpylon/sessions/session_subject.py b/mtpylon/sessions/session_subject.py
new file mode 100644
index 0000000..9dae38b
--- /dev/null
+++ b/mtpylon/sessions/session_subject.py
@@ -0,0 +1,51 @@
+# -*- coding: utf-8 -*-
+from asyncio import gather
+from typing import Callable
+
+from mtpylon.types import long
+from .session_subject_protocol import SessionSubjectProtocol
+from .session_storage_protocol import SessionStorageProtocol
+from .session_oberver_protocol import SessionObserverProtocol
+from .session_event import SessionEvent
+
+
+class SessionSubject(SessionSubjectProtocol):
+
+    def __init__(
+        self,
+        session_storage_factory: Callable[[], SessionStorageProtocol]
+    ):
+        self.session_storage = session_storage_factory()
+        self.observers = []
+
+    def subscribe(self, observer: SessionObserverProtocol):
+        self.observers.append(observer)
+
+    def unsubscribe(self, observer: SessionObserverProtocol):
+        self.observers.remove(observer)
+
+    async def _notify(self, event: SessionEvent):
+        await gather(*[
+            observer.update(event)
+            for observer in self.observers
+        ])
+
+    async def create_session(self, session_id: long):
+        await self.session_storage.create_session(session_id)
+        event = SessionEvent(
+            type='created',
+            session_id=session_id
+        )
+        await self._notify(event)
+
+    async def has_session(self, session_id: long) -> bool:
+        return await self.session_storage.has_session(session_id)
+
+    async def destroy_session(self, session_id: long):
+        if await self.session_storage.has_session(session_id):
+            await self.session_storage.destroy_session(session_id)
+            event = SessionEvent(
+                type='destroyed',
+                session_id=session_id
+            )
+            await self._notify(event)
diff --git a/mtpylon/sessions/session_subject_protocol.py b/mtpylon/sessions/session_subject_protocol.py
new file mode 100644
index 0000000..a1c9c3c
--- /dev/null
+++ b/mtpylon/sessions/session_subject_protocol.py
@@ -0,0 +1,21 @@
+# -*- coding: utf-8 -*-
+from typing import List
+
+from .session_storage_protocol import SessionStorageProtocol
+from .session_oberver_protocol import SessionObserverProtocol
+
+
+class SessionSubjectProtocol(SessionStorageProtocol):
+
+    session_storage: SessionStorageProtocol
+
+    observers: List[SessionObserverProtocol]
+
+    def subscribe(self, observer: SessionObserverProtocol):  # pragma: nocover
+        ...
+
+    def unsubscribe(
+        self,
+        observer: SessionObserverProtocol
+    ):  # pragma: nocover
+        ...


Test patch:
diff --git a/tests/test_aiohandlers/test_websockets.py b/tests/test_aiohandlers/test_websockets.py
index 6a0df33..80dba06 100644
--- a/tests/test_aiohandlers/test_websockets.py
+++ b/tests/test_aiohandlers/test_websockets.py
@@ -9,6 +9,7 @@
 from mtpylon.crypto import AuthKeyManager
 from mtpylon.dh_prime_generators.single_prime import generate
 from mtpylon.salts import ServerSaltManager
+from mtpylon.sessions import SessionSubject, InMemorySessionStorage
 
 from tests.helpers import hexstr_to_bytes
 from tests.simple_manager import manager
@@ -119,6 +120,32 @@ async def test_no_server_salt_manager(self):
             assert conn.closed
 
 
+class WsHandlerNoSessionSubject(AioHTTPTestCase):
+
+    async def get_application(self) -> web.Application:
+        ws_handler = create_websocket_handler(schema)
+
+        app = web.Application()
+        app['rsa_manager'] = manager
+        app['auth_key_manager'] = AuthKeyManager()
+        app['dh_prime_generator'] = generate()
+        app['server_salt_manager'] = ServerSaltManager()
+
+        app.router.add_get('/ws', ws_handler)
+
+        return app
+
+    @unittest_run_loop
+    async def test_no_session_subject(self):
+        logger = MagicMock()
+
+        with patch('mtpylon.aiohandlers.websockets.logger', logger):
+            async with self.client.ws_connect('/ws') as conn:
+                assert logger.error.called
+
+            assert conn.closed
+
+
 class WsHandlerTestCase(AioHTTPTestCase):
 
     async def get_application(self) -> web.Application:
@@ -129,6 +156,9 @@ async def get_application(self) -> web.Application:
         app['rsa_manager'] = manager
         app['dh_prime_generator'] = generate()
         app['server_salt_manager'] = ServerSaltManager()
+        app['session_subject'] = SessionSubject(
+            lambda: InMemorySessionStorage()
+        )
         app.router.add_get('/ws', ws_handler)
 
         return app
@@ -190,7 +220,7 @@ async def test_valid_transport_tag_passed(self):
                 await conn.send_bytes(good_header)
                 assert not conn.closed
 
-        assert not logger.error.called
+        assert not logger.error.called, logger.error.call_args[0]
         assert MessageHandler.called
 
     @unittest_run_loop
@@ -225,7 +255,7 @@ async def test_valid_transport_tag_message(self):
                 await conn.send_bytes(clients_message)
                 assert not conn.closed
 
-        assert not logger.error.called
+        assert not logger.error.called, logger.error.call_args[0]
         assert MessageHandler.called
         assert MessageSender.called
         assert message_entity.handle.called
diff --git a/tests/test_sessions/__init__.py b/tests/test_sessions/__init__.py
new file mode 100644
index 0000000..40a96af
--- /dev/null
+++ b/tests/test_sessions/__init__.py
@@ -0,0 +1,1 @@
+# -*- coding: utf-8 -*-
diff --git a/tests/test_sessions/test_in_memory_session_storage.py b/tests/test_sessions/test_in_memory_session_storage.py
new file mode 100644
index 0000000..6437dd0
--- /dev/null
+++ b/tests/test_sessions/test_in_memory_session_storage.py
@@ -0,0 +1,56 @@
+# -*- coding: utf-8 -*-
+import pytest
+
+from mtpylon import long
+from mtpylon.sessions import InMemorySessionStorage
+
+
+@pytest.fixture
+def session_storage():
+    return InMemorySessionStorage()
+
+
+@pytest.mark.asyncio
+async def test_create_session(session_storage):
+    session_id = long(123123)
+
+    await session_storage.create_session(session_id)
+
+    assert len(session_storage._storage) == 1
+    assert session_id in session_storage._storage
+
+
+@pytest.mark.asyncio
+async def test_has_not_got_session(session_storage):
+    session_id = long(123123)
+
+    assert not await session_storage.has_session(session_id)
+
+
+@pytest.mark.asyncio
+async def test_has_got_session(session_storage):
+    session_id = long(123123)
+
+    await session_storage.create_session(session_id)
+
+    assert await session_storage.has_session(session_id)
+
+
+@pytest.mark.asyncio
+async def test_destroy_session(session_storage):
+    session_id = long(123123)
+
+    await session_storage.create_session(session_id)
+
+    await session_storage.destroy_session(session_id)
+
+    assert not await session_storage.has_session(session_id)
+
+
+@pytest.mark.asyncio
+async def test_destroy_session_no_id(session_storage):
+    session_id = long(123123)
+
+    await session_storage.destroy_session(session_id)
+
+    assert not await session_storage.has_session(session_id)
diff --git a/tests/test_sessions/test_session_subject.py b/tests/test_sessions/test_session_subject.py
new file mode 100644
index 0000000..e5b81be
--- /dev/null
+++ b/tests/test_sessions/test_session_subject.py
@@ -0,0 +1,106 @@
+# -*- coding: utf-8 -*-
+from unittest.mock import MagicMock, AsyncMock
+
+import pytest
+
+from mtpylon.types import long
+from mtpylon.sessions import (
+    SessionEvent,
+    SessionSubject,
+    InMemorySessionStorage
+)
+
+
+@pytest.fixture
+def session_subject():
+    return SessionSubject(lambda: InMemorySessionStorage())
+
+
+@pytest.mark.asyncio
+async def test_create_session(session_subject):
+    observer1 = MagicMock(update=AsyncMock())
+    observer2 = MagicMock(update=AsyncMock())
+
+    session_subject.subscribe(observer1)
+    session_subject.subscribe(observer2)
+
+    session_id = long(123123)
+
+    await session_subject.create_session(session_id)
+
+    created_event = SessionEvent(
+        type='created',
+        session_id=session_id,
+    )
+
+    observer1.update.assert_awaited()
+    assert observer1.update.await_args[0][0] == created_event
+
+    observer2.update.assert_awaited()
+    assert observer2.update.await_args[0][0] == created_event
+
+
+@pytest.mark.asyncio
+async def test_destroy_session(session_subject):
+    session_id = long(123123)
+    await session_subject.create_session(session_id)
+
+    observer1 = MagicMock(update=AsyncMock())
+    observer2 = MagicMock(update=AsyncMock())
+    session_subject.subscribe(observer1)
+    session_subject.subscribe(observer2)
+
+    await session_subject.destroy_session(session_id)
+
+    destroyed_event = SessionEvent(
+        type='destroyed',
+        session_id=session_id,
+    )
+
+    observer1.update.assert_awaited()
+    assert observer1.update.await_args[0][0] == destroyed_event
+
+    observer2.update.assert_awaited()
+    assert observer2.update.await_args[0][0] == destroyed_event
+
+
+@pytest.mark.asyncio
+async def test_destory_session_no_session(session_subject):
+    session_id = long(123123)
+
+    observer1 = MagicMock(update=AsyncMock())
+
+    session_subject.subscribe(observer1)
+    await session_subject.destroy_session(session_id)
+
+    observer1.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_unsubscribe_observer(session_subject):
+    session1_id = long(123123)
+    session2_id = long(234234)
+
+    observer1 = MagicMock(update=AsyncMock())
+    observer2 = MagicMock(update=AsyncMock())
+
+    session_subject.subscribe(observer1)
+    session_subject.subscribe(observer2)
+
+    await session_subject.create_session(session1_id)
+
+    session_subject.unsubscribe(observer1)
+
+    await session_subject.create_session(session2_id)
+
+    assert observer1.update.await_count == 1
+    assert observer2.update.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_has_got_session(session_subject):
+    session_id = long(123123)
+
+    await session_subject.create_session(session_id)
+
+    assert await session_subject.has_session(session_id)