Subscribe Unsubscribe message handlers (#7)

Author: nkiryanov

src/handlers/tests/conftest.py renamed to src/app/fixtures.py

@@ -1,14 +1,13 @@
 import pytest
 
 from app.testing import MockedWebSocketServerProtocol
-from storage import SubscriptionStorage
 
 
 @pytest.fixture
-def storage():
-    return SubscriptionStorage()
+def ws():
+    return MockedWebSocketServerProtocol()
 
 
 @pytest.fixture
-def ws():
+def ya_ws():
     return MockedWebSocketServerProtocol()

src/conftest.py

@@ -2,6 +2,11 @@
 
 from app.conf import Settings
 
+pytest_plugins = [
+    "app.fixtures",
+    "storage.fixtures",
+]
+
 
 @pytest.fixture(autouse=True)
 def settings(mocker):

src/handlers/__init__.py

@@ -1,5 +1,5 @@
-from handlers.websocket_auth_message_handler import WebsocketAuthMessageHandler
+from handlers.message_handler import WebSocketMessageHandler
 
 __all__ = [
-    "WebsocketAuthMessageHandler",
+    "WebSocketMessageHandler",
 ]

src/handlers/message_handler.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from typing import Coroutine, Any, Callable
+
+from websockets import WebSocketServerProtocol
+
+from a12n.jwk_client import AsyncJWKClient
+from a12n.jwk_client import AsyncJWKClientException
+from app import conf
+from handlers.dto import AuthMessage
+from handlers.dto import SubscribeMessage
+from handlers.dto import UnsubscribeMessage
+from handlers.dto import SuccessResponseMessage
+from handlers.exceptions import WebsocketMessageException
+from storage.exceptions import StorageOperationException
+from storage.storage_updaters import StorageWebSocketRegister
+from storage import SubscriptionStorage
+from storage.storage_updaters import StorageUserSubscriber
+from storage.storage_updaters import StorageUserUnsubscriber
+from handlers.dto import IncomingMessage
+
+AsyncMessageHandler = Callable[[WebSocketServerProtocol, Any], Coroutine[Any, Any, SuccessResponseMessage]]
+
+
+@dataclass
+class WebSocketMessageHandler:
+    storage: SubscriptionStorage
+
+    def __post_init__(self) -> None:
+        settings = conf.get_app_settings()
+        self.jwk_client = AsyncJWKClient(jwks_url=settings.AUTH_JWKS_URL, supported_signing_algorithms=settings.AUTH_SUPPORTED_SIGNING_ALGORITHMS)
+
+        self.message_handlers: dict[str, AsyncMessageHandler] = {
+            "Authenticate": self.handle_auth_message,
+            "Subscribe": self.handle_subscribe_message,
+            "Unsubscribe": self.handle_unsubscribe_message,
+        }
+
+    async def handle_message(self, websocket: WebSocketServerProtocol, message: IncomingMessage) -> SuccessResponseMessage:
+        return await self.message_handlers[message.message_type](websocket, message)
+
+    async def handle_auth_message(self, websocket: WebSocketServerProtocol, message: AuthMessage) -> SuccessResponseMessage:
+        try:
+            validated_token = await self.jwk_client.decode(message.params.token)
+            StorageWebSocketRegister(storage=self.storage, websocket=websocket, validated_token=validated_token)()
+        except (AsyncJWKClientException, StorageOperationException) as exc:
+            raise WebsocketMessageException(str(exc), message)
+
+        return SuccessResponseMessage.model_construct(incoming_message=message)
+
+    async def handle_subscribe_message(self, websocket: WebSocketServerProtocol, message: SubscribeMessage) -> SuccessResponseMessage:
+        StorageUserSubscriber(storage=self.storage, websocket=websocket, event=message.params.event)()
+        return SuccessResponseMessage.model_construct(incoming_message=message)
+
+    async def handle_unsubscribe_message(self, websocket: WebSocketServerProtocol, message: UnsubscribeMessage) -> SuccessResponseMessage:
+        StorageUserUnsubscriber(storage=self.storage, websocket=websocket, event=message.params.event)()
+        return SuccessResponseMessage.model_construct(incoming_message=message)

src/handlers/tests/message_handler/conftest.py

@@ -0,0 +1,36 @@
+import pytest
+
+from handlers import WebSocketMessageHandler
+from handlers.dto import AuthMessage, SubscribeMessage, UnsubscribeMessage
+
+
+@pytest.fixture(autouse=True)
+def settings(settings):
+    settings.AUTH_JWKS_URL = "https://auth.clowns.com/auth/realms/clowns-realm/protocol/openid-connect/certs"
+    settings.AUTH_SUPPORTED_SIGNING_ALGORITHMS = ["RS256"]
+    return settings
+
+
+@pytest.fixture
+def force_token_on_validation(mocker, valid_token):
+    return mocker.patch("a12n.jwk_client.AsyncJWKClient.decode", return_value=valid_token)
+
+
+@pytest.fixture
+def message_handler(storage):
+    return WebSocketMessageHandler(storage=storage)
+
+
+@pytest.fixture
+def auth_message():
+    return AuthMessage(message_id=23, message_type="Authenticate", params={"token": "some-valid-token-value"})
+
+
+@pytest.fixture
+def subscribe_message():
+    return SubscribeMessage(message_id=24, message_type="Subscribe", params={"event": "channel1"})
+
+
+@pytest.fixture
+def unsubscribe_message():
+    return UnsubscribeMessage(message_id=25, message_type="Unsubscribe", params={"event": "channel1"})

src/handlers/tests/message_handler/tests_auth_message_handler.py

@@ -0,0 +1,67 @@
+import pytest
+
+from a12n.jwk_client import AsyncJWKClientException
+from app.types import DecodedValidToken
+from handlers.dto import SuccessResponseMessage
+from handlers.exceptions import WebsocketMessageException
+from handlers import WebSocketMessageHandler
+from storage.storage_updaters import StorageWebSocketRegister
+
+pytestmark = [
+    pytest.mark.usefixtures("force_token_on_validation"),
+]
+
+
+@pytest.fixture
+def ya_user_decoded_valid_token():
+    return DecodedValidToken(sub="ya_user", exp="4852128170")
+
+
+@pytest.fixture
+def auth_handler(message_handler: WebSocketMessageHandler, ws):
+    return lambda auth_message: message_handler.handle_auth_message(ws, auth_message)
+
+
+async def test_auth_handler_response_on_correct_auth_message(auth_handler, auth_message):
+    auth_response = await auth_handler(auth_message)
+
+    assert isinstance(auth_response, SuccessResponseMessage)
+    assert auth_response.message_type == "SuccessResponse"
+    assert auth_response.incoming_message == auth_message
+
+
+async def test_auth_handler_register_websocket_in_storage(auth_handler, ws, auth_message, mocker, storage, valid_token):
+    spy_websocket_register = mocker.spy(StorageWebSocketRegister, "__call__")
+
+    await auth_handler(auth_message)
+
+    assert storage.is_websocket_registered(ws) is True
+    spy_websocket_register.assert_called_once()
+    called_service = spy_websocket_register.call_args.args[0]
+    assert called_service.storage == storage
+    assert called_service.websocket == ws
+    assert called_service.validated_token == valid_token
+
+
+async def test_auth_handler_raise_if_user_send_token_for_different_user(auth_handler, auth_message, storage, ws, register_ws, ya_user_decoded_valid_token):
+    register_ws(ws, ya_user_decoded_valid_token)
+
+    with pytest.raises(WebsocketMessageException) as exc_info:
+        await auth_handler(auth_message)  # send valid user1 token while connection registered with ya_user
+
+    raised_exception = exc_info.value
+    assert raised_exception.error_detail == "The user has different public id"
+    assert raised_exception.incoming_message == auth_message
+    assert storage.is_websocket_registered(ws) is True, "The existed connection should not be touched"
+
+
+async def test_auth_handler_raise_if_user_try_to_auth_with_expired_token(auth_handler, ws, auth_message, force_token_on_validation, storage):
+    force_token_on_validation.side_effect = AsyncJWKClientException("The token is expired")
+
+    with pytest.raises(WebsocketMessageException) as exc_info:
+        await auth_handler(auth_message)
+
+    raised_exception = exc_info.value
+    assert raised_exception.error_detail == "The token is expired"
+    assert raised_exception.incoming_message == auth_message
+    assert storage.is_websocket_registered(ws) is False, "The ws should not be added to registered websockets"

src/handlers/tests/message_handler/tests_message_handler_common.py

@@ -0,0 +1,37 @@
+import pytest
+from handlers.message_handler import WebSocketMessageHandler
+
+
+@pytest.fixture
+def get_message_handler(storage):
+    return lambda: WebSocketMessageHandler(storage)
+
+
+def test_message_handler_jwk_client_settings(message_handler):
+    assert message_handler.jwk_client.jwks_url == "https://auth.clowns.com/auth/realms/clowns-realm/protocol/openid-connect/certs"
+    assert message_handler.jwk_client.supported_signing_algorithms == ["RS256"]
+
+
+@pytest.mark.usefixtures("force_token_on_validation")
+async def test_message_handler_call_auth_handler_on_auth_message(get_message_handler, auth_message, mocker, ws):
+    spy_auth_handler = mocker.spy(WebSocketMessageHandler, "handle_auth_message")
+
+    await get_message_handler().handle_message(ws, auth_message)
+
+    spy_auth_handler.assert_awaited_once()
+
+
+async def test_message_handler_call_subscribe_handler_on_subscribe_message(get_message_handler, subscribe_message, mocker, ws_registered):
+    spy_subscribe_handler = mocker.spy(WebSocketMessageHandler, "handle_subscribe_message")
+
+    await get_message_handler().handle_message(ws_registered, subscribe_message)
+
+    spy_subscribe_handler.assert_awaited_once()
+
+
+async def test_message_handler_call_unsubscribe_handler_on_unsubscribe_message(get_message_handler, unsubscribe_message, mocker, ws_subscribed):
+    spy_unsubscribe_handler = mocker.spy(WebSocketMessageHandler, "handle_unsubscribe_message")
+
+    await get_message_handler().handle_message(ws_subscribed, unsubscribe_message)
+
+    spy_unsubscribe_handler.assert_awaited_once()

src/handlers/tests/message_handler/tests_subscirbe_message_handler.py

@@ -0,0 +1,27 @@
+import pytest
+from handlers.message_handler import WebSocketMessageHandler
+from storage.storage_updaters import StorageUserSubscriber
+
+
+@pytest.fixture
+def subscribe_handler(message_handler: WebSocketMessageHandler, ws_registered):
+    return lambda subscribe_message: message_handler.handle_subscribe_message(ws_registered, subscribe_message)
+
+
+async def test_subscribe_handler_return_success_response(subscribe_handler, subscribe_message):
+    subscribe_response = await subscribe_handler(subscribe_message)
+
+    assert subscribe_response.message_type == "SuccessResponse"
+    assert subscribe_response.incoming_message == subscribe_message
+
+
+async def test_subscribe_handler_call_storage_subscriber_under_the_hood(subscribe_handler, subscribe_message, mocker, storage, ws_registered):
+    spy_storage_subscriber = mocker.spy(StorageUserSubscriber, "__call__")
+
+    await subscribe_handler(subscribe_message)
+
+    spy_storage_subscriber.assert_called_once()
+    called_service = spy_storage_subscriber.call_args.args[0]
+    assert called_service.storage == storage
+    assert called_service.websocket == ws_registered
+    assert called_service.event == subscribe_message.params.event

src/handlers/tests/message_handler/tests_unsubscirbe_message_handler.py

@@ -0,0 +1,27 @@
+import pytest
+from handlers.message_handler import WebSocketMessageHandler
+from storage.storage_updaters import StorageUserUnsubscriber
+
+
+@pytest.fixture
+def unsubscribe_handler(message_handler: WebSocketMessageHandler, ws_subscribed):
+    return lambda unsubscribe_message: message_handler.handle_unsubscribe_message(ws_subscribed, unsubscribe_message)
+
+
+async def test_unsubscribe_handler_return_success_response(unsubscribe_handler, unsubscribe_message):
+    unsubscribe_response = await unsubscribe_handler(unsubscribe_message)
+
+    assert unsubscribe_response.message_type == "SuccessResponse"
+    assert unsubscribe_response.incoming_message == unsubscribe_message
+
+
+async def test_unsubscribe_handler_call_storage_unsubscriber_under_the_hood(unsubscribe_handler, unsubscribe_message, mocker, storage, ws_subscribed):
+    spy_storage_unsubscriber = mocker.spy(StorageUserUnsubscriber, "__call__")
+
+    await unsubscribe_handler(unsubscribe_message)
+
+    spy_storage_unsubscriber.assert_called_once()
+    called_service = spy_storage_unsubscriber.call_args.args[0]
+    assert called_service.storage == storage
+    assert called_service.websocket == ws_subscribed
+    assert called_service.event == unsubscribe_message.params.event