fixed failing test

Author: eadwinCode

ellar/common/interfaces/__init__.py

@@ -15,6 +15,7 @@
 from .interceptor_consumer import IInterceptorsConsumer
 from .middleware import IEllarMiddleware
 from .module import IModuleSetup
+from .operation import IWebSocketConnectionAttributes
 from .response_model import IResponseModel
 from .templating import IModuleTemplateLoader, ITemplateRenderingService
 from .versioning import IAPIVersioning, IAPIVersioningResolver
@@ -43,4 +44,5 @@
     "IIdentitySchemes",
     "IApplicationReady",
     "ITemplateRenderingService",
+    "IWebSocketConnectionAttributes",
 ]

ellar/common/interfaces/operation.py

@@ -0,0 +1,23 @@
+import typing as t
+
+
+class IWebSocketConnectionAttributes(t.Protocol):
+    """
+    Interface for WebSocket connection attributes.
+    """
+
+    def connect(self, websocket_handler: t.Callable) -> t.Callable:
+        """
+        Register the connect handler to a websocket handler.
+
+        :param websocket_handler: The websocket handler to register the connect handler to.
+        :return: The connect handler.
+        """
+
+    def disconnect(self, websocket_handler: t.Callable) -> t.Callable:
+        """
+        Register the disconnect handler to a websocket handler.
+
+        :param websocket_handler: The websocket handler to register the disconnect handler to.
+        :return: The disconnect handler.
+        """

ellar/common/operations/base.py

@@ -1,6 +1,5 @@
 import functools
 import typing as t
-from functools import partial
 from types import FunctionType
 
 from ellar.common.constants import (
@@ -15,55 +14,18 @@
     ROUTE_OPERATION_PARAMETERS,
     TRACE,
 )
+from ellar.common.interfaces.operation import IWebSocketConnectionAttributes
 from ellar.reflect import ensure_target
 
 from .schema import RouteParameters, WsRouteParameters
 
 
-class _WebSocketConnectionAttributes:
-    """
-    This class is used to add connection attributes to a websocket handler.
-    """
-
-    __slots__ = ("_original_func", "_func")
-
-    def __init__(self, func: t.Callable) -> None:
-        self._original_func = func
-        self._func: t.Optional[t.Callable] = None
-
-    def __call__(
-        self,
-        path: str = "/",
-        *,
-        name: t.Optional[str] = None,
-        encoding: t.Optional[str] = "json",
-        use_extra_handler: bool = False,
-        extra_handler_type: t.Optional[t.Type] = None,
-    ) -> t.Callable:
-        if self._func is None:  # pragma: no cover
-            raise Exception("Something went wrong")
-
-        res = self._func(
-            path=path,
-            name=name,
-            encoding=encoding,
-            use_extra_handler=use_extra_handler,
-            extra_handler_type=extra_handler_type,
-        )
-        return t.cast(t.Callable, res)
-
-    def __get__(self, instance: t.Any, owner: t.Any) -> t.Callable:
-        self._func = functools.partial(self._original_func, instance)
-        return self
-
-    @classmethod
+def _websocket_connection_attributes(
+    func: t.Callable,
+) -> IWebSocketConnectionAttributes:
     def _advance_function(
-        cls, websocket_handler: t.Callable, handler_name: str
+        websocket_handler: t.Callable, handler_name: str
     ) -> t.Callable:
-        """
-        This method is used to register the connection attributes to a websocket handler.
-        """
-
         def _wrap(connect_handler: t.Callable) -> t.Callable:
             if not (
                 callable(websocket_handler) and type(websocket_handler) is FunctionType
@@ -89,17 +51,10 @@ def _wrap(connect_handler: t.Callable) -> t.Callable:
 
         return _wrap
 
-    def connect(self, f: t.Callable) -> t.Callable:
-        """
-        This method is used to register the connect handler to a websocket handler.
-        """
-        return self._advance_function(f, "on_connect")
+    func.connect = functools.partial(_advance_function, handler_name="on_connect")  # type: ignore[attr-defined]
+    func.disconnect = functools.partial(_advance_function, handler_name="on_disconnect")  # type: ignore[attr-defined]
 
-    def disconnect(self, f: t.Callable) -> t.Callable:
-        """
-        This method is used to register the disconnect handler to a websocket handler.
-        """
-        return self._advance_function(f, "on_disconnect")
+    return t.cast(IWebSocketConnectionAttributes, func)
 
 
 class OperationDefinitions:
@@ -182,7 +137,7 @@ def get(
         ] = None,
     ) -> t.Callable:
         methods = [GET]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -202,7 +157,7 @@ def post(
         ] = None,
     ) -> t.Callable:
         methods = [POST]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -222,7 +177,7 @@ def put(
         ] = None,
     ) -> t.Callable:
         methods = [PUT]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -242,7 +197,7 @@ def patch(
         ] = None,
     ) -> t.Callable:
         methods = [PATCH]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -262,7 +217,7 @@ def delete(
         ] = None,
     ) -> t.Callable:
         methods = [DELETE]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -282,7 +237,7 @@ def head(
         ] = None,
     ) -> t.Callable:
         methods = [HEAD]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -302,7 +257,7 @@ def options(
         ] = None,
     ) -> t.Callable:
         methods = [OPTIONS]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -322,7 +277,7 @@ def trace(
         ] = None,
     ) -> t.Callable:
         methods = [TRACE]
-        endpoint_parameter_partial = partial(
+        endpoint_parameter_partial = functools.partial(
             RouteParameters,
             name=name,
             methods=methods,
@@ -356,7 +311,7 @@ def _decorator(endpoint_handler: t.Callable) -> t.Callable:
 
         return _decorator
 
-    @_WebSocketConnectionAttributes
+    @_websocket_connection_attributes
     def ws_route(
         self,
         path: str = "/",

ellar/core/router_builders/controller.py

@@ -32,7 +32,9 @@ def process_controller_routes(controller: t.Type[ControllerBase]) -> t.List[Base
     for _, item in get_functions_with_tag(
         controller, tag=constants.OPERATION_ENDPOINT_KEY
     ):
-        parameters = item.__dict__[constants.ROUTE_OPERATION_PARAMETERS]
+        parameters = item.__dict__.get(constants.ROUTE_OPERATION_PARAMETERS)
+        if parameters is None:
+            print("Something is not right")
         operation: t.Union[ControllerRouteOperation, ControllerWebsocketRouteOperation]
 
         if not isinstance(parameters, list):