Add functionality to mimic a websocket client against an endpoint

.gitignore

@@ -17,3 +17,5 @@ build/*
 .DS_Store
 dist/*
 pip-wheel-metadata/
+.venv
+.vscode

sanic_testing/reusable.py

@@ -1,15 +1,16 @@
 import asyncio
+import typing
 from functools import partial
 from random import randint
-from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, Tuple
 
 import httpx
 from sanic import Sanic
 from sanic.application.state import ApplicationServerInfo
 from sanic.log import logger
 from sanic.request import Request
-from websockets.legacy.client import connect
+
+from sanic_testing.websocket import websocket_proxy
 
 from .testing import HOST, PORT, TestingResponse
 
@@ -157,11 +158,7 @@ async def _local_request(self, method, url, *args, **kwargs):
         raw_cookies = kwargs.pop("raw_cookies", None)
 
         if method == "websocket":
-            ws_proxy = SimpleNamespace()
-            async with connect(url, *args, **kwargs) as websocket:
-                ws_proxy.ws = websocket
-                ws_proxy.opened = True
-            return ws_proxy
+            return await websocket_proxy(url, *args, **kwargs)
         else:
             session = self._session
 
@@ -225,5 +222,13 @@ def options(self, *args, **kwargs):
     def head(self, *args, **kwargs):
         return self._sanic_endpoint_test("head", *args, **kwargs)
 
-    def websocket(self, *args, **kwargs):
+    def websocket(
+        self,
+        *args,
+        mimic: typing.Optional[
+            typing.Callable[..., typing.Coroutine[None, None, typing.Any]]
+        ] = None,
+        **kwargs,
+    ):
+        kwargs["mimic"] = mimic
         return self._sanic_endpoint_test("websocket", *args, **kwargs)

sanic_testing/testing.py

@@ -4,7 +4,6 @@
 from json import JSONDecodeError
 from socket import AF_INET6, SOCK_STREAM, socket
 from string import ascii_lowercase
-from types import SimpleNamespace
 
 import httpx
 from sanic import Sanic  # type: ignore
@@ -13,7 +12,8 @@
 from sanic.log import logger  # type: ignore
 from sanic.request import Request  # type: ignore
 from sanic.response import text  # type: ignore
-from websockets.legacy.client import connect
+
+from sanic_testing.websocket import websocket_proxy
 
 ASGI_HOST = "mockserver"
 ASGI_PORT = 1234
@@ -94,11 +94,7 @@ async def _local_request(self, method: str, url: str, *args, **kwargs):
                 kwargs["follow_redirects"] = allow_redirects
 
         if method == "websocket":
-            ws_proxy = SimpleNamespace()
-            async with connect(url, *args, **kwargs) as websocket:
-                ws_proxy.ws = websocket
-                ws_proxy.opened = True
-            return ws_proxy
+            return await websocket_proxy(url, *args, **kwargs)
         else:
             async with self.get_new_session(**session_kwargs) as session:
 
@@ -305,7 +301,15 @@ def options(self, *args, **kwargs):
     def head(self, *args, **kwargs):
         return self._sanic_endpoint_test("head", *args, **kwargs)
 
-    def websocket(self, *args, **kwargs):
+    def websocket(
+        self,
+        *args,
+        mimic: typing.Optional[
+            typing.Callable[..., typing.Coroutine[None, None, typing.Any]]
+        ] = None,
+        **kwargs,
+    ):
+        kwargs["mimic"] = mimic
         return self._sanic_endpoint_test("websocket", *args, **kwargs)
 
 
@@ -397,7 +401,21 @@ async def _ws_receive(cls):
     async def _ws_send(cls, message):
         pass
 
-    async def websocket(self, uri, subprotocols=None, *args, **kwargs):
+    async def websocket(
+        self,
+        uri,
+        subprotocols=None,
+        *args,
+        mimic: typing.Optional[
+            typing.Callable[..., typing.Coroutine[None, None, typing.Any]]
+        ] = None,
+        **kwargs,
+    ):
+        if mimic:
+            raise RuntimeError(
+                "SanicASGITestClient does not currently support the mimic "
+                "keyword argument. Please use SanicTestClient instead."
+            )
         scheme = "ws"
         path = uri
         root_path = f"{scheme}://{ASGI_HOST}:{ASGI_PORT}"

sanic_testing/websocket.py

@@ -0,0 +1,49 @@
+import typing
+
+from websockets.exceptions import ConnectionClosedOK
+from websockets.legacy.client import connect
+
+
+class WebsocketProxy:
+    def __init__(self, ws):
+        self.ws = ws
+        self.opened = True
+        self.client_received: typing.List[str] = []
+        self.client_sent: typing.List[str] = []
+
+    @property
+    def server_received(self):
+        return self.client_sent
+
+    @property
+    def server_sent(self):
+        return self.client_received
+
+
+async def websocket_proxy(url, *args, **kwargs) -> WebsocketProxy:
+    mimic = kwargs.pop("mimic", None)
+    async with connect(url, *args, **kwargs) as websocket:
+        ws_proxy = WebsocketProxy(websocket)
+
+        if mimic:
+            do_send = websocket.send
+            do_recv = websocket.recv
+
+            async def send(data):
+                ws_proxy.client_sent.append(data)
+                await do_send(data)
+
+            async def recv():
+                message = await do_recv()
+                ws_proxy.client_received.append(message)
+
+            websocket.send = send  # type: ignore
+            websocket.recv = recv  # type: ignore
+
+            try:
+                await mimic(websocket)
+            except ConnectionClosedOK:
+                pass
+            else:
+                await websocket.send("")
+    return ws_proxy

tests/test_test_client.py

@@ -1,7 +1,9 @@
 import asyncio
 
 import pytest
+from sanic import Sanic, Websocket
 from sanic.request import Request
+from websockets.client import WebSocketClientProtocol
 
 
 @pytest.mark.parametrize(
@@ -16,7 +18,7 @@ def test_basic_test_client(app, method):
     assert response.content_type == "text/plain; charset=utf-8"
 
 
-def test_websocket_route(app):
+def test_websocket_route_basic(app):
     ev = asyncio.Event()
 
     @app.websocket("/ws")
@@ -30,6 +32,35 @@ async def handler(request, ws):
     assert ev.is_set()
 
 
+def test_websocket_route_queue(app: Sanic):
+    async def client_mimic(websocket: WebSocketClientProtocol):
+        await websocket.send("foo")
+        await websocket.recv()
+
+    @app.websocket("/ws")
+    async def handler(request, ws: Websocket):
+        while True:
+            await ws.send("hello!")
+            if not await ws.recv():
+                break
+
+    _, response = app.test_client.websocket("/ws", mimic=client_mimic)
+    assert response.server_sent == ["hello!"]
+    assert response.server_received == ["foo", ""]
+
+
+def test_websocket_client_mimic_failed(app: Sanic):
+    @app.websocket("/ws")
+    async def handler(request, ws: Websocket):
+        pass
+
+    async def client_mimic(websocket: WebSocketClientProtocol):
+        raise Exception("Should fails")
+
+    with pytest.raises(Exception, match="Should fails"):
+        app.test_client.websocket("/ws", mimic=client_mimic)
+
+
 def test_listeners(app):
     listeners = []
     available = (