Repository: sanic-org/sanic
Pull request #2841 (2023-10-20 21:48:39)

Add websocket.handler.before/after/exception signals
Adding ability to implement middleware functionality on WebSocket handlers by signals that receive `websocket` parameter.

The flow is:
```
http.handler.before (request) -> websocket.handler.before (request, websocket) -> handler
 ____________________________________________________________________________________/
|
| -> clean exit: websocket.handler.after (request, websocket)
| -> except Exception: websocket.handler.exception (request, websocket, exception)
| -> cancelled/closed: no signal
                             |
                      any of the above -> http.handler.after(request)
```


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/sanic-org/sanic/pull/2841?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`a5a9658`)](https://app.codecov.io/gh/sanic-org/sanic/commit/a5a9658896984ddad484e168d4cb5c96e589fbad?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org) 88.387% compared to head [(`50cec6b`)](https://app.codecov.io/gh/sanic-org/sanic/pull/2841?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org) 88.452%.

<details><summary>Additional details and impacted files</summary>


```diff
@@              Coverage Diff              @@
##              main     #2841       +/-   ##
=============================================
+ Coverage   88.387%   88.452%   +0.065%     
=============================================
  Files           92        92               
  Lines         7156      7162        +6     
  Branches      1228      1228               
=============================================
+ Hits          6325      6335       +10     
+ Misses         575       573        -2     
+ Partials       256       254        -2     
```


| [Files](https://app.codecov.io/gh/sanic-org/sanic/pull/2841?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org) | Coverage Δ | |
|---|---|---|
| [sanic/app.py](https://app.codecov.io/gh/sanic-org/sanic/pull/2841?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org#diff-c2FuaWMvYXBwLnB5) | `88.803% <100.000%> (+0.051%)` | :arrow_up: |
| [sanic/signals.py](https://app.codecov.io/gh/sanic-org/sanic/pull/2841?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org#diff-c2FuaWMvc2lnbmFscy5weQ==) | `91.724% <100.000%> (+0.174%)` | :arrow_up: |

... and [2 files with indirect coverage changes](https://app.codecov.io/gh/sanic-org/sanic/pull/2841/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/sanic-org/sanic/pull/2841?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sanic-org).

Tronic: <img src=https://zi.fi/ws-signals.png>
Diagram created by ChatGPT based on the ASCII sketch in OP.

Patch:
diff --git a/sanic/app.py b/sanic/app.py
index 452954ca..12786a18 100644
--- a/sanic/app.py
+++ b/sanic/app.py
@@ -1369,6 +1369,12 @@ class Sanic(
             protocol = request.transport.get_protocol()
             ws = await protocol.websocket_handshake(request, subprotocols)
 
+        await self.dispatch(
+            "websocket.handler.before",
+            inline=True,
+            context={"request": request, "websocket": ws},
+            fail_not_found=False,
+        )
         # schedule the application handler
         # its future is kept in self.websocket_tasks in case it
         # needs to be cancelled due to the server being stopped
@@ -1377,10 +1383,24 @@ class Sanic(
         cancelled = False
         try:
             await fut
+            await self.dispatch(
+                "websocket.handler.after",
+                inline=True,
+                context={"request": request, "websocket": ws},
+                reverse=True,
+                fail_not_found=False,
+            )
         except (CancelledError, ConnectionClosed):  # type: ignore
             cancelled = True
         except Exception as e:
             self.error_handler.log(request, e)
+            await self.dispatch(
+                "websocket.handler.exception",
+                inline=True,
+                context={"request": request, "websocket": ws, "exception": e},
+                reverse=True,
+                fail_not_found=False,
+            )
         finally:
             self.websocket_tasks.remove(fut)
             if cancelled:
diff --git a/sanic/signals.py b/sanic/signals.py
index fe252c12..b5069068 100644
--- a/sanic/signals.py
+++ b/sanic/signals.py
@@ -38,6 +38,9 @@ class Event(Enum):
     HTTP_LIFECYCLE_SEND = "http.lifecycle.send"
     HTTP_MIDDLEWARE_AFTER = "http.middleware.after"
     HTTP_MIDDLEWARE_BEFORE = "http.middleware.before"
+    WEBSOCKET_HANDLER_AFTER = "websocket.handler.after"
+    WEBSOCKET_HANDLER_BEFORE = "websocket.handler.before"
+    WEBSOCKET_HANDLER_EXCEPTION = "websocket.handler.exception"
 
 
 RESERVED_NAMESPACES = {
@@ -65,6 +68,11 @@ RESERVED_NAMESPACES = {
         Event.HTTP_MIDDLEWARE_AFTER.value,
         Event.HTTP_MIDDLEWARE_BEFORE.value,
     ),
+    "websocket": {
+        Event.WEBSOCKET_HANDLER_AFTER.value,
+        Event.WEBSOCKET_HANDLER_BEFORE.value,
+        Event.WEBSOCKET_HANDLER_EXCEPTION.value,
+    },
 }
 
 


Test patch:
diff --git a/tests/test_ws_handlers.py b/tests/test_ws_handlers.py
index 5205090d..c70dbeed 100644
--- a/tests/test_ws_handlers.py
+++ b/tests/test_ws_handlers.py
@@ -54,3 +54,65 @@ def test_ws_handler_async_for(
     )
     assert ws_proxy.client_sent == ["test 1", "test 2", ""]
     assert ws_proxy.client_received == ["test 1", "test 2"]
+
+
+def signalapp(app):
+    @app.signal("websocket.handler.before")
+    async def ws_before(request: Request, websocket: Websocket):
+        app.ctx.seq.append("before")
+        print("before")
+        await websocket.send("before: " + await websocket.recv())
+        print("before2")
+
+    @app.signal("websocket.handler.after")
+    async def ws_after(request: Request, websocket: Websocket):
+        app.ctx.seq.append("after")
+        await websocket.send("after: " + await websocket.recv())
+        await websocket.recv()
+
+    @app.signal("websocket.handler.exception")
+    async def ws_exception(
+        request: Request, websocket: Websocket, exception: Exception
+    ):
+        app.ctx.seq.append("exception")
+        await websocket.send(f"exception: {exception}")
+        await websocket.recv()
+
+    @app.websocket("/ws")
+    async def ws_handler(request: Request, ws: Websocket):
+        app.ctx.seq.append("ws")
+
+    @app.websocket("/wserror")
+    async def ws_error(request: Request, ws: Websocket):
+        print("wserr")
+        app.ctx.seq.append("wserror")
+        raise Exception(await ws.recv())
+        print("wserr2")
+
+
+def test_ws_signals(
+    app: Sanic,
+    simple_ws_mimic_client: MimicClientType,
+):
+    signalapp(app)
+
+    app.ctx.seq = []
+    _, ws_proxy = app.test_client.websocket(
+        "/ws", mimic=simple_ws_mimic_client
+    )
+    assert ws_proxy.client_received == ["before: test 1", "after: test 2"]
+    assert app.ctx.seq == ["before", "ws", "after"]
+
+
+def test_ws_signals_exception(
+    app: Sanic,
+    simple_ws_mimic_client: MimicClientType,
+):
+    signalapp(app)
+
+    app.ctx.seq = []
+    _, ws_proxy = app.test_client.websocket(
+        "/wserror", mimic=simple_ws_mimic_client
+    )
+    assert ws_proxy.client_received == ["before: test 1", "exception: test 2"]
+    assert app.ctx.seq == ["before", "wserror", "exception"]