From 2b4660bd113462acf767c465d0cafa9cc2949b16 Mon Sep 17 00:00:00 2001 From: Adam Hopkins Date: Sun, 24 Jul 2022 23:52:10 +0300 Subject: [PATCH 1/3] Initial implementation of websocket mimic --- sanic_testing/testing.py | 33 +++++++++++++++++++++++++++++++++ tests/test_test_client.py | 20 +++++++++++++++++++- 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/sanic_testing/testing.py b/sanic_testing/testing.py index 6dd7fe2..dc84e78 100644 --- a/sanic_testing/testing.py +++ b/sanic_testing/testing.py @@ -1,4 +1,5 @@ import typing +from asyncio import sleep from functools import partial from ipaddress import IPv6Address, ip_address from json import JSONDecodeError @@ -13,6 +14,7 @@ from sanic.log import logger # type: ignore from sanic.request import Request # type: ignore from sanic.response import text # type: ignore +from websockets.exceptions import ConnectionClosedOK from websockets.legacy.client import connect ASGI_HOST = "mockserver" @@ -95,9 +97,40 @@ async def _local_request(self, method: str, url: str, *args, **kwargs): if method == "websocket": ws_proxy = SimpleNamespace() + mimic = kwargs.pop("mimic", None) async with connect(url, *args, **kwargs) as websocket: ws_proxy.ws = websocket ws_proxy.opened = True + ws_proxy.received = [] + ws_proxy.sent = [] + + if mimic: + do_send = websocket.send + do_recv = websocket.recv + + async def send(data): + ws_proxy.received.append(data) + await do_send(data) + + async def recv(): + message = await do_recv() + ws_proxy.sent.append(message) + + websocket.send = send # type: ignore + websocket.recv = recv # type: ignore + + async def do_mimic(): + try: + await mimic(websocket) + except ConnectionClosedOK: + ... + else: + await websocket.send("") + + task = self.app.loop.create_task(do_mimic()) + + while not task.done(): + await sleep(0.1) return ws_proxy else: async with self.get_new_session(**session_kwargs) as session: diff --git a/tests/test_test_client.py b/tests/test_test_client.py index 4601200..ab0693d 100644 --- a/tests/test_test_client.py +++ b/tests/test_test_client.py @@ -1,6 +1,7 @@ import asyncio import pytest +from sanic import Websocket from sanic.request import Request @@ -16,7 +17,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 +31,23 @@ async def handler(request, ws): assert ev.is_set() +def test_websocket_route_queue(app): + async def client_mimic(websocket): + 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.sent == ["hello!"] + assert response.received == ["foo", ""] + + def test_listeners(app): listeners = [] available = ( From b61d7097c191895ad2d42cd04d38cb612bad5ba9 Mon Sep 17 00:00:00 2001 From: Zhiwei Date: Mon, 25 Jul 2022 16:09:33 -0500 Subject: [PATCH 2/3] Refactor, Cleanup, and Test Case for WebSocket Minic Client Implementation (#51) --- .gitignore | 2 ++ sanic_testing/testing.py | 17 ++++------------- tests/test_test_client.py | 21 +++++++++++++++++---- 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/.gitignore b/.gitignore index 91cac37..327761e 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,5 @@ build/* .DS_Store dist/* pip-wheel-metadata/ +.venv +.vscode diff --git a/sanic_testing/testing.py b/sanic_testing/testing.py index dc84e78..2f1e3b8 100644 --- a/sanic_testing/testing.py +++ b/sanic_testing/testing.py @@ -1,5 +1,4 @@ import typing -from asyncio import sleep from functools import partial from ipaddress import IPv6Address, ip_address from json import JSONDecodeError @@ -119,18 +118,10 @@ async def recv(): websocket.send = send # type: ignore websocket.recv = recv # type: ignore - async def do_mimic(): - try: - await mimic(websocket) - except ConnectionClosedOK: - ... - else: - await websocket.send("") - - task = self.app.loop.create_task(do_mimic()) - - while not task.done(): - await sleep(0.1) + try: + await mimic(websocket) + except ConnectionClosedOK: + pass return ws_proxy else: async with self.get_new_session(**session_kwargs) as session: diff --git a/tests/test_test_client.py b/tests/test_test_client.py index ab0693d..d5c8948 100644 --- a/tests/test_test_client.py +++ b/tests/test_test_client.py @@ -1,8 +1,9 @@ import asyncio import pytest -from sanic import Websocket +from sanic import Sanic, Websocket from sanic.request import Request +from websockets.client import WebSocketClientProtocol @pytest.mark.parametrize( @@ -31,8 +32,8 @@ async def handler(request, ws): assert ev.is_set() -def test_websocket_route_queue(app): - async def client_mimic(websocket): +def test_websocket_route_queue(app: Sanic): + async def client_mimic(websocket: WebSocketClientProtocol): await websocket.send("foo") await websocket.recv() @@ -45,7 +46,19 @@ async def handler(request, ws: Websocket): _, response = app.test_client.websocket("/ws", mimic=client_mimic) assert response.sent == ["hello!"] - assert response.received == ["foo", ""] + assert response.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"): + _, response = app.test_client.websocket("/ws", mimic=client_mimic) def test_listeners(app): From bb3de5d5ca7cd525b0ba55bef80b48df43abfd50 Mon Sep 17 00:00:00 2001 From: Adam Hopkins Date: Tue, 26 Jul 2022 00:36:06 +0300 Subject: [PATCH 3/3] Generalize websocket proxy --- sanic_testing/reusable.py | 21 ++++++++----- sanic_testing/testing.py | 60 +++++++++++++++++--------------------- sanic_testing/websocket.py | 49 +++++++++++++++++++++++++++++++ tests/test_test_client.py | 6 ++-- 4 files changed, 92 insertions(+), 44 deletions(-) create mode 100644 sanic_testing/websocket.py diff --git a/sanic_testing/reusable.py b/sanic_testing/reusable.py index 12d10c5..160bff0 100644 --- a/sanic_testing/reusable.py +++ b/sanic_testing/reusable.py @@ -1,7 +1,7 @@ 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 @@ -9,7 +9,8 @@ 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) diff --git a/sanic_testing/testing.py b/sanic_testing/testing.py index 2f1e3b8..3260e4c 100644 --- a/sanic_testing/testing.py +++ b/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,8 +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.exceptions import ConnectionClosedOK -from websockets.legacy.client import connect + +from sanic_testing.websocket import websocket_proxy ASGI_HOST = "mockserver" ASGI_PORT = 1234 @@ -95,34 +94,7 @@ async def _local_request(self, method: str, url: str, *args, **kwargs): kwargs["follow_redirects"] = allow_redirects if method == "websocket": - ws_proxy = SimpleNamespace() - mimic = kwargs.pop("mimic", None) - async with connect(url, *args, **kwargs) as websocket: - ws_proxy.ws = websocket - ws_proxy.opened = True - ws_proxy.received = [] - ws_proxy.sent = [] - - if mimic: - do_send = websocket.send - do_recv = websocket.recv - - async def send(data): - ws_proxy.received.append(data) - await do_send(data) - - async def recv(): - message = await do_recv() - ws_proxy.sent.append(message) - - websocket.send = send # type: ignore - websocket.recv = recv # type: ignore - - try: - await mimic(websocket) - except ConnectionClosedOK: - pass - return ws_proxy + return await websocket_proxy(url, *args, **kwargs) else: async with self.get_new_session(**session_kwargs) as session: @@ -329,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) @@ -421,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}" diff --git a/sanic_testing/websocket.py b/sanic_testing/websocket.py new file mode 100644 index 0000000..8c6138a --- /dev/null +++ b/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 diff --git a/tests/test_test_client.py b/tests/test_test_client.py index d5c8948..b2e72f0 100644 --- a/tests/test_test_client.py +++ b/tests/test_test_client.py @@ -45,8 +45,8 @@ async def handler(request, ws: Websocket): break _, response = app.test_client.websocket("/ws", mimic=client_mimic) - assert response.sent == ["hello!"] - assert response.received == ["foo"] + assert response.server_sent == ["hello!"] + assert response.server_received == ["foo", ""] def test_websocket_client_mimic_failed(app: Sanic): @@ -58,7 +58,7 @@ async def client_mimic(websocket: WebSocketClientProtocol): raise Exception("Should fails") with pytest.raises(Exception, match="Should fails"): - _, response = app.test_client.websocket("/ws", mimic=client_mimic) + app.test_client.websocket("/ws", mimic=client_mimic) def test_listeners(app):