Added testing

Author: eadwinCode

docs/overview/custom_decorators.md

@@ -46,12 +46,14 @@ from ellar.core import WebSocket
 def sample_endpoint_ws(self, websocket: WebSocket, data_schema: str = WsBody()):
     pass
 
-@sample_endpoint_ws.connect
+@ws_route.connect(sample_endpoint_ws)
 async def on_connect(self, websocket: WebSocket):
+    # Called when there is a connection to `sample_endpoint_ws`
     await websocket.accept()
 
-@sample_endpoint_ws.disconnect
-async def on_connect(self, websocket: WebSocket, code: int):
+@ws_route.disconnect(sample_endpoint_ws)
+async def on_disconnect(self, websocket: WebSocket, code: int):
+    # Called when there is a disconnect from `sample_endpoint_ws`
     await websocket.close(code)
 ```
 

tests/test_socket_io/__init__.py

@@ -0,0 +1,130 @@
+import typing as t
+from functools import wraps
+
+from starlette import status
+from starlette.exceptions import WebSocketException
+
+from ellar.common import Guards, Header, Query, WsBody, extra_args
+from ellar.core.connection import HTTPConnection
+from ellar.core.guard import APIKeyHeader
+from ellar.core.params import ExtraEndpointArg
+from ellar.core.serializer import Serializer
+from ellar.di import injectable
+from ellar.socket_io import (
+    WebSocketGateway,
+    WsResponse,
+    on_connected,
+    on_disconnected,
+    subscribe_message,
+)
+from ellar.socket_io.model import GatewayBase
+
+
+def add_extra_args(func):
+    # EXTRA ARGS SETUP
+    query1 = ExtraEndpointArg(name="query1", annotation=str, default_value=Query())
+    query2 = ExtraEndpointArg(
+        name="query2", annotation=str
+    )  # will default to Query during computation
+
+    extra_args(query1, query2)(func)
+
+    @wraps(func)
+    def _wrapper(*args, **kwargs):
+        # RESOLVING EXTRA ARGS
+        # All extra args must be resolved before calling route function
+        # else extra argument will be pushed to the route function
+        resolved_query1 = query1.resolve(kwargs)
+        resolved_query2 = query2.resolve(kwargs)
+
+        return func(*args, **kwargs, query2=resolved_query2, query1=resolved_query1)
+
+    return _wrapper
+
+
+@injectable()
+class HeaderGuard(APIKeyHeader):
+    parameter_name = "x-auth-key"
+
+    async def authenticate(
+        self, connection: HTTPConnection, key: t.Optional[t.Any]
+    ) -> t.Optional[t.Any]:
+        if key == "supersecret":
+            return key
+
+
+class MessageData(Serializer):
+    data: t.Any
+
+
+class MessageRoom(Serializer):
+    room: str
+
+
+class MessageChatRoom(Serializer):
+    room: str
+    data: t.Any
+
+
+@WebSocketGateway(path="/ws", async_mode="asgi", cors_allowed_origins="*")
+class EventGateway:
+    @subscribe_message("my_event")
+    async def my_event(self, message: MessageData = WsBody()):
+        return WsResponse("my_response", {"data": message.data}, room=self.context.sid)
+
+    @subscribe_message
+    async def my_broadcast_event(self, message: MessageData = WsBody()):
+        await self.context.server.emit("my_response", {"data": message.data})
+
+    @on_connected()
+    async def connect(self):
+        await self.context.server.emit(
+            "my_response", {"data": "Connected", "count": 0}, room=self.context.sid
+        )
+
+    @on_disconnected()
+    async def disconnect(self):
+        print("Client disconnected")
+
+
+@Guards(HeaderGuard)
+@WebSocketGateway(path="/ws-guard")
+class GatewayWithGuards(GatewayBase):
+    @subscribe_message("my_event")
+    async def my_event(self, message: MessageData = WsBody()):
+        return WsResponse(
+            "my_response",
+            {
+                "data": message.data,
+                "auth-key": self.context.switch_to_http_connection().get_client().user,
+            },
+        )
+
+    @subscribe_message("my_event_header")
+    async def my_event_header(
+        self,
+        data: str = WsBody(..., embed=True),
+        x_auth_key: str = Header(alias="x-auth-key"),
+    ):
+        return WsResponse("my_response", {"data": data, "x_auth_key": x_auth_key})
+
+
+@WebSocketGateway(path="/ws-others")
+class GatewayOthers(GatewayBase):
+    @subscribe_message("my_event")
+    async def my_event(self, message: MessageData = WsBody()):
+        raise Exception("I dont have anything to run.")
+
+    @subscribe_message("my_event_raise")
+    async def my_event_raise(self, message: MessageData = WsBody()):
+        raise WebSocketException(
+            code=status.WS_1009_MESSAGE_TOO_BIG, reason="Message is too big"
+        )
+
+    @subscribe_message("extra_args")
+    @add_extra_args
+    def extra_args_handler(self, query1: str, query2: str):
+        raise WebSocketException(
+            code=status.WS_1009_MESSAGE_TOO_BIG,
+            reason={"query1": query1, "query2": query2},
+        )

tests/test_socket_io/sample.py

@@ -0,0 +1,25 @@
+from ellar.reflect import reflect
+from ellar.socket_io import on_connected, on_disconnected
+from ellar.socket_io.constants import (
+    CONNECTION_EVENT,
+    DISCONNECT_EVENT,
+    MESSAGE_MAPPING_METADATA,
+)
+
+
+def test_on_connected_decorator_works():
+    @on_connected()
+    def sample_function():
+        pass
+
+    assert getattr(sample_function, MESSAGE_MAPPING_METADATA)
+    assert reflect.get_metadata_or_raise_exception(CONNECTION_EVENT, sample_function)
+
+
+def test_on_disconnected_decorator_works():
+    @on_disconnected()
+    def sample_function():
+        pass
+
+    assert getattr(sample_function, MESSAGE_MAPPING_METADATA)
+    assert reflect.get_metadata_or_raise_exception(DISCONNECT_EVENT, sample_function)

tests/test_socket_io/test_decorators/test_events.py

@@ -0,0 +1,128 @@
+import pytest
+
+from ellar.common import Guards
+from ellar.constants import CONTROLLER_CLASS_KEY, GUARDS_KEY
+from ellar.core.guard import HttpBearerAuth
+from ellar.helper import get_name
+from ellar.reflect import reflect
+from ellar.socket_io import WebSocketGateway, subscribe_message
+from ellar.socket_io.constants import (
+    GATEWAY_MESSAGE_HANDLER_KEY,
+    GATEWAY_METADATA,
+    GATEWAY_OPTIONS,
+    GATEWAY_WATERMARK,
+    MESSAGE_METADATA,
+)
+from ellar.socket_io.model import GatewayBase, GatewayType
+
+
+@WebSocketGateway(path="/ws", namespace="/some-namespace")
+@Guards(HttpBearerAuth)
+class SampleWithoutGateway:
+    pass
+
+
+@WebSocketGateway(path="/ws", namespace="/some-namespace")
+class SampleWithGateway(GatewayBase):
+    pass
+
+
+@WebSocketGateway
+class SampleMarkAsGateway:
+    pass
+
+
+@pytest.mark.parametrize(
+    "gateway, watermark, options",
+    [
+        (
+            SampleWithoutGateway,
+            True,
+            {
+                "async_mode": "asgi",
+                "cors_allowed_origins": "*",
+                "namespace": "/some-namespace",
+            },
+        ),
+        (
+            SampleWithGateway,
+            False,
+            {
+                "async_mode": "asgi",
+                "cors_allowed_origins": "*",
+                "namespace": "/some-namespace",
+            },
+        ),
+        (
+            SampleMarkAsGateway,
+            True,
+            {"async_mode": "asgi", "cors_allowed_origins": "*"},
+        ),
+    ],
+)
+def test_websocket_gateway_works_without_gateway(gateway, watermark, options):
+
+    assert isinstance(gateway, GatewayType)
+    assert hasattr(gateway, "__GATEWAY_WATERMARK__") is watermark
+
+    assert reflect.get_metadata_or_raise_exception(GATEWAY_WATERMARK, gateway) is True
+    assert reflect.get_metadata_or_raise_exception(GATEWAY_OPTIONS, gateway) == options
+
+    for key in GATEWAY_METADATA.keys:
+        reflect.get_metadata_or_raise_exception(key, gateway)
+
+
+def test_gateway_guard():
+    assert reflect.get_metadata_or_raise_exception(
+        GUARDS_KEY, SampleWithoutGateway
+    ) == [HttpBearerAuth]
+
+
+def test_sub_message_building_works():
+    @WebSocketGateway(path="/ws", namespace="/some-namespace")
+    class SampleAGateway(GatewayBase):
+        @subscribe_message
+        def a_message(self):
+            pass
+
+    message_handlers = reflect.get_metadata_or_raise_exception(
+        GATEWAY_MESSAGE_HANDLER_KEY, SampleAGateway
+    )
+    assert len(message_handlers) == 1
+    assert get_name(message_handlers[0]) == "a_message"
+    message = reflect.get_metadata_or_raise_exception(
+        MESSAGE_METADATA, message_handlers[0]
+    )
+    assert message == "a_message"
+    assert (
+        reflect.get_metadata_or_raise_exception(
+            CONTROLLER_CLASS_KEY, SampleAGateway().a_message
+        )
+        is SampleAGateway
+    )
+
+
+def test_sub_message_building_fails():
+    with pytest.raises(Exception) as ex:
+
+        @WebSocketGateway(path="/ws", namespace="/some-namespace")
+        class SampleBGateway(GatewayBase):
+            @subscribe_message
+            @reflect.metadata(CONTROLLER_CLASS_KEY, "b_message")
+            def b_message(self):
+                pass
+
+    assert (
+        "SampleBGateway Gateway message handler tried to be processed more than once"
+        in str(ex.value)
+    )
+
+
+def test_cant_use_gateway_decorator_on_function():
+    with pytest.raises(Exception) as ex:
+
+        @WebSocketGateway(path="/ws", namespace="/some-namespace")
+        def sample_c_gateway():
+            pass
+
+    assert "WebSocketGateway is a class decorator" in str(ex.value)

tests/test_socket_io/test_decorators/test_gateway.py

@@ -0,0 +1,25 @@
+from ellar.reflect import reflect
+from ellar.socket_io import subscribe_message
+from ellar.socket_io.constants import MESSAGE_MAPPING_METADATA, MESSAGE_METADATA
+
+
+def test_subscribe_message_works():
+    @subscribe_message("sample")
+    def sample_function():
+        pass
+
+    assert getattr(sample_function, MESSAGE_MAPPING_METADATA)
+    assert (
+        reflect.get_metadata_or_raise_exception(MESSAGE_METADATA, sample_function)
+        == "sample"
+    )
+
+    @subscribe_message
+    def sample_function2():
+        pass
+
+    assert getattr(sample_function2, MESSAGE_MAPPING_METADATA)
+    assert (
+        reflect.get_metadata_or_raise_exception(MESSAGE_METADATA, sample_function2)
+        == "sample_function2"
+    )