Added more test on websocket

Author: eadwinCode

ellar/core/routing/operation_definitions.py

@@ -263,7 +263,7 @@ def ws_route(
         path: str = "/",
         *,
         name: str = None,
-        encoding: str = "json",
+        encoding: t.Optional[str] = "json",
         use_extra_handler: bool = False,
         extra_handler_type: t.Optional[t.Type] = None,
     ) -> t.Callable[[TCallable], t.Union[TCallable, TWebsocketOperation]]:

ellar/core/routing/websocket/handler.py

@@ -65,9 +65,13 @@ async def dispatch(
                 elif message["type"] == "websocket.disconnect":
                     close_code = int(message.get("code", status.WS_1000_NORMAL_CLOSURE))
                     break
+        except WebSocketException as wexc:
+            await websocket.close(code=wexc.code)
+            raise RuntimeError(wexc.reason)
         except Exception as exc:
             close_code = status.WS_1011_INTERNAL_ERROR
             raise exc
+
         finally:
             await self.execute_on_disconnect(context=context, close_code=close_code)
 

ellar/core/schema.py

@@ -84,7 +84,7 @@ class WsRouteParameters(BaseModel):
     path: str
     name: t.Optional[str] = None
     endpoint: t.Callable
-    encoding: str = Field("json")
+    encoding: t.Optional[str] = Field("json")
     use_extra_handler: bool = Field(False)
     extra_handler_type: t.Optional[t.Type["WebSocketExtraHandler"]] = None
 
@@ -94,7 +94,7 @@ def validate_endpoint(cls, value: t.Any):
 
     @validator("encoding")
     def validate_encoding(cls, value: t.Any):
-        if value not in ["json", "text", "bytes"]:
+        if value not in ["json", "text", "bytes", None]:
             raise ValueError(
                 f"Encoding type not supported. Once [json | text | bytes]. Received: {value}"
             )

tests/test_websocket_handler.py

@@ -1,13 +1,13 @@
 import pytest
-from starlette.websockets import WebSocket
+from starlette.websockets import WebSocket, WebSocketState
 
-from ellar.common import ModuleRouter, WsBody
+from ellar.common import Controller, ModuleRouter, WsBody, ws_route
 from ellar.core import TestClientFactory
 from ellar.core.exceptions import ImproperConfiguration
 
 from .schema import Item
 
-router = ModuleRouter("/")
+router = ModuleRouter("/router")
 
 
 @router.ws_route("/ws-with-handler", use_extra_handler=True)
@@ -24,6 +24,12 @@ async def websocket_with_handler_connect(websocket: WebSocket):
     await websocket.accept()
 
 
+@websocket_with_handler.disconnect
+async def websocket_with_handler_connect(websocket: WebSocket, code: int):
+    # await websocket.close(code=code)
+    assert websocket.client_state == WebSocketState.DISCONNECTED
+
+
 @router.ws_route("/ws")
 async def websocket_without_handler(websocket: WebSocket, query: str):
     assert query == "my-query"
@@ -33,12 +39,45 @@ async def websocket_without_handler(websocket: WebSocket, query: str):
     await websocket.close()
 
 
-tm = TestClientFactory.create_test_module(routers=[router])
+@Controller("/controller")
+class WebsocketController:
+    @ws_route("/ws-with-handler", use_extra_handler=True)
+    async def websocket_with_handler_c(
+        self, websocket: WebSocket, query: str, data: Item = WsBody()
+    ):
+        assert query == "my-query"
+        await websocket.send_json(data.dict())
+        await websocket.close()
+
+    @websocket_with_handler_c.connect
+    async def websocket_with_handler_connect(self, websocket: WebSocket):
+        await websocket.accept()
+
+    @websocket_with_handler_c.disconnect
+    async def websocket_with_handler_connect(self, websocket: WebSocket, code: int):
+        # await websocket.close(code=code)
+        assert websocket.client_state == WebSocketState.DISCONNECTED
+
+    @ws_route("/ws")
+    async def websocket_without_handler_c(self, websocket: WebSocket, query: str):
+        assert query == "my-query"
+        await websocket.accept()
+        message = await websocket.receive_text()
+        await websocket.send_json({"message": f"Thanks. {message}"})
+        await websocket.close()
+
+
+tm = TestClientFactory.create_test_module(
+    routers=[router], controllers=(WebsocketController,)
+)
 client = tm.get_client()
 
 
-def test_websocket_with_handler_works():
-    with client.websocket_connect("/ws-with-handler?query=my-query") as session:
+@pytest.mark.parametrize("prefix", ["/router", "/controller"])
+def test_websocket_with_handler_works(prefix):
+    with client.websocket_connect(
+        f"{prefix}/ws-with-handler?query=my-query"
+    ) as session:
         session.send_json(Item(name="Ellar", price=23.34, tax=1.2).dict())
         message = session.receive_json()
         assert message == {
@@ -49,9 +88,12 @@ def test_websocket_with_handler_works():
         }
 
 
-def test_websocket_with_handler_fails_for_invalid_input():
+@pytest.mark.parametrize("prefix", ["/router", "/controller"])
+def test_websocket_with_handler_fails_for_invalid_input(prefix):
     with pytest.raises(Exception):
-        with client.websocket_connect("/ws-with-handler?query=my-query") as session:
+        with client.websocket_connect(
+            f"{prefix}/ws-with-handler?query=my-query"
+        ) as session:
             session.send_json({"framework": "Ellar is awesome"})
             message = session.receive_json()
     assert message == {
@@ -71,9 +113,10 @@ def test_websocket_with_handler_fails_for_invalid_input():
     }
 
 
-def test_websocket_with_handler_fails_for_missing_route_parameter():
+@pytest.mark.parametrize("prefix", ["/router", "/controller"])
+def test_websocket_with_handler_fails_for_missing_route_parameter(prefix):
     with pytest.raises(Exception):
-        with client.websocket_connect("/ws-with-handler") as session:
+        with client.websocket_connect(f"{prefix}/ws-with-handler") as session:
             session.send_json(Item(name="Ellar", price=23.34, tax=1.2).dict())
             message = session.receive_json()
     assert message == {
@@ -88,8 +131,9 @@ def test_websocket_with_handler_fails_for_missing_route_parameter():
     }
 
 
-def test_plain_websocket_route():
-    with client.websocket_connect("/ws?query=my-query") as websocket:
+@pytest.mark.parametrize("prefix", ["/router", "/controller"])
+def test_plain_websocket_route(prefix):
+    with client.websocket_connect(f"{prefix}/ws?query=my-query") as websocket:
         websocket.send_text("Ellar")
         message = websocket.receive_json()
         assert message == {"message": "Thanks. Ellar"}
@@ -106,3 +150,139 @@ async def websocket_with_handler(
             websocket: WebSocket, query: str, data: Item = WsBody()
         ):
             pass
+
+
+def test_websocket_endpoint_on_connect():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True)
+        async def ws(self, websocket: WebSocket):
+            pass
+
+        @ws.connect
+        async def on_connect(self, websocket):
+            assert websocket["subprotocols"] == ["soap", "wamp"]
+            await websocket.accept(subprotocol="wamp")
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+    with _client.websocket_connect("/ws/", subprotocols=["soap", "wamp"]) as websocket:
+        assert websocket.accepted_subprotocol == "wamp"
+
+
+def test_websocket_endpoint_on_receive_bytes():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True, encoding="bytes")
+        async def ws(self, websocket: WebSocket, data: bytes = WsBody()):
+            await websocket.send_bytes(b"Message bytes was: " + data)
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+    with _client.websocket_connect("/ws/") as websocket:
+        websocket.send_bytes(b"Hello, world!")
+        _bytes = websocket.receive_bytes()
+        assert _bytes == b"Message bytes was: Hello, world!"
+
+    with pytest.raises(RuntimeError):
+        with _client.websocket_connect("/ws/") as websocket:
+            websocket.send_text("Hello world")
+
+
+def test_websocket_endpoint_on_receive_json():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True, encoding="json")
+        async def ws(self, websocket: WebSocket, data=WsBody()):
+            await websocket.send_json({"message": data})
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+
+    with _client.websocket_connect("/ws/") as websocket:
+        websocket.send_json({"hello": "world"})
+        data = websocket.receive_json()
+        assert data == {"message": {"hello": "world"}}
+
+    with pytest.raises(RuntimeError):
+        with _client.websocket_connect("/ws/") as websocket:
+            websocket.send_text("Hello world")
+
+
+def test_websocket_endpoint_on_receive_json_binary():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True, encoding="json")
+        async def ws(self, websocket: WebSocket, data=WsBody()):
+            await websocket.send_json({"message": data}, mode="binary")
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+
+    with _client.websocket_connect("/ws/") as websocket:
+        websocket.send_json({"hello": "world"}, mode="binary")
+        data = websocket.receive_json(mode="binary")
+        assert data == {"message": {"hello": "world"}}
+
+
+def test_websocket_endpoint_on_receive_text():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True, encoding="text")
+        async def ws(self, websocket: WebSocket, data: str = WsBody()):
+            await websocket.send_text(f"Message text was: {data}")
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+
+    with _client.websocket_connect("/ws/") as websocket:
+        websocket.send_text("Hello, world!")
+        _text = websocket.receive_text()
+        assert _text == "Message text was: Hello, world!"
+
+    with pytest.raises(RuntimeError):
+        with _client.websocket_connect("/ws/") as websocket:
+            websocket.send_bytes(b"Hello world")
+
+
+def test_websocket_endpoint_on_default():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True, encoding=None)
+        async def ws(self, websocket: WebSocket, data: str = WsBody()):
+            await websocket.send_text(f"Message text was: {data}")
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+
+    with _client.websocket_connect("/ws/") as websocket:
+        websocket.send_text("Hello, world!")
+        _text = websocket.receive_text()
+        assert _text == "Message text was: Hello, world!"
+
+
+def test_websocket_endpoint_on_disconnect():
+    @Controller("/ws")
+    class WebSocketSample:
+        @ws_route(use_extra_handler=True, encoding=None)
+        async def ws(self, websocket: WebSocket, data: str = WsBody()):
+            await websocket.send_text(f"Message text was: {data}")
+
+        @ws.disconnect
+        async def on_disconnect(self, websocket: WebSocket, close_code):
+            assert close_code == 1001
+            await websocket.close(code=close_code)
+
+    _client = TestClientFactory.create_test_module(
+        controllers=(WebSocketSample,)
+    ).get_client()
+
+    with _client.websocket_connect("/ws/") as websocket:
+        websocket.send_text("Hello, world!")
+        websocket.close(code=1001)