Allow user to redefine websocket.accept in BaseWSEndpoint

dolamroth · Mar 7, 2024 · 5aed29b · 5aed29b
1 parent fc4711b
commit 5aed29b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/starlette_web/common/ws/base_endpoint.py b/starlette_web/common/ws/base_endpoint.py
@@ -61,10 +61,13 @@ async def on_connect(self, websocket: WebSocket) -> None:
         del websocket.state.db_session
 
         if permitted:
-            await websocket.accept()
+            await self.accept(websocket)
         else:
             raise WebSocketDisconnect(code=3000, reason=reason)
 
+    async def accept(self, websocket: WebSocket) -> None:
+        await websocket.accept()
+
     async def on_receive(self, websocket: WebSocket, data: Any) -> None:
         cleaned_data = self._validate(data)
         task_id = get_random_string(50)