From be711b52fd6a9f5eca47d92fe47e888033228585 Mon Sep 17 00:00:00 2001 From: Jake Holland Date: Sun, 12 Jun 2022 17:59:46 +0100 Subject: [PATCH 1/4] Added support for websocket immediate closure handling --- async_asgi_testclient/tests/test_testing.py | 39 +++++++++++++++++++++ async_asgi_testclient/websocket.py | 8 +++-- 2 files changed, 45 insertions(+), 2 deletions(-) diff --git a/async_asgi_testclient/tests/test_testing.py b/async_asgi_testclient/tests/test_testing.py index c48a9a0..594c82a 100644 --- a/async_asgi_testclient/tests/test_testing.py +++ b/async_asgi_testclient/tests/test_testing.py @@ -1,3 +1,7 @@ +import ast + +import starlette.status + from async_asgi_testclient import TestClient from http.cookies import SimpleCookie from json import dumps @@ -108,6 +112,11 @@ async def on_receive(self, websocket, data): else: await websocket.send_text(f"Message text was: {data}") + @app.websocket_route("/ws-reject") + async def websocket_reject(websocket): + # Send immediate close message to the client, using non default 100 code to test return of correct code + await websocket.close(starlette.status.WS_1003_UNSUPPORTED_DATA) + @app.route("/") async def homepage(request): return Response("full response") @@ -450,6 +459,36 @@ async def test_ws_connect_custom_scheme(starlette_app): assert msg.startswith("wss://") +@pytest.mark.asyncio +async def test_ws_endpoint_with_immediate_rejection(starlette_app): + async with TestClient(starlette_app, timeout=0.1) as client: + try: + async with client.websocket_connect("/ws-reject") as ws: + pass + except Exception as e: + thrown_exception = e + + assert ast.literal_eval(str(thrown_exception)) == { + "type": "websocket.close", + "code": starlette.status.WS_1003_UNSUPPORTED_DATA + } + + +@pytest.mark.asyncio +async def test_invalid_ws_endpoint(starlette_app): + async with TestClient(starlette_app, timeout=0.1) as client: + try: + async with client.websocket_connect("/invalid") as ws: + pass + except Exception as e: + thrown_exception = e + + assert ast.literal_eval(str(thrown_exception)) == { + "type": "websocket.close", + "code": starlette.status.WS_1000_NORMAL_CLOSURE + } + + @pytest.mark.asyncio async def test_request_stream(starlette_app): from starlette.responses import StreamingResponse diff --git a/async_asgi_testclient/websocket.py b/async_asgi_testclient/websocket.py index 4f75024..e895cd2 100644 --- a/async_asgi_testclient/websocket.py +++ b/async_asgi_testclient/websocket.py @@ -87,7 +87,10 @@ async def receive_json(self): return json.loads(data) async def _receive(self): - return await receive(self.output_queue) + try: + return await receive(self.output_queue) + except: + pass def __aiter__(self): return self @@ -131,4 +134,5 @@ async def connect(self): await self._send({"type": "websocket.connect"}) msg = await self._receive() - assert msg["type"] == "websocket.accept" + if msg["type"] != "websocket.accept": + raise Exception(msg) From 017ac376b1fa9453b7485157eec1456873113241 Mon Sep 17 00:00:00 2001 From: Jake Holland Date: Sun, 12 Jun 2022 18:05:22 +0100 Subject: [PATCH 2/4] Added tests for quart websockets - Same tests as the starlette app ones, just using the quart app --- async_asgi_testclient/tests/test_testing.py | 98 +++++++++++++++++++++ 1 file changed, 98 insertions(+) diff --git a/async_asgi_testclient/tests/test_testing.py b/async_asgi_testclient/tests/test_testing.py index 594c82a..ed92f14 100644 --- a/async_asgi_testclient/tests/test_testing.py +++ b/async_asgi_testclient/tests/test_testing.py @@ -1,6 +1,7 @@ import ast import starlette.status +from quart import websocket from async_asgi_testclient import TestClient from http.cookies import SimpleCookie @@ -84,6 +85,20 @@ async def echoheaders(): async def test_query(): return Response(request.query_string) + @app.websocket("/ws") + async def websocket_endpoint(): + data = await websocket.receive() + if data == "cookies": + await websocket.send(dumps(websocket.cookies)) + elif data == "url": + await websocket.send(str(websocket.url)) + else: + await websocket.send(f"Message text was: {data}") + + @app.websocket("/ws-reject") + async def websocket_reject(): + await websocket.close(code=starlette.status.WS_1003_UNSUPPORTED_DATA, reason="some reason") + yield app @@ -489,6 +504,89 @@ async def test_invalid_ws_endpoint(starlette_app): } +@pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") +async def test_quart_ws_endpoint(quart_app): + async with TestClient(quart_app, timeout=0.1) as client: + async with client.websocket_connect("/ws") as ws: + await ws.send_text("hi!") + msg = await ws.receive_text() + assert msg == "Message text was: hi!" + + +@pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") +async def test_quart_ws_endpoint_cookies(quart_app): + async with TestClient(quart_app, timeout=0.1) as client: + async with client.websocket_connect("/ws", cookies={"session": "abc"}) as ws: + await ws.send_text("cookies") + msg = await ws.receive_json() + assert msg == {"session": "abc"} + + +@pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") +async def test_quart_ws_connect_inherits_test_client_cookies(quart_app): + client = TestClient(quart_app, use_cookies=True, timeout=0.1) + client.cookie_jar = SimpleCookie({"session": "abc"}) + async with client: + async with client.websocket_connect("/ws") as ws: + await ws.send_text("cookies") + msg = await ws.receive_text() + assert msg == '{"session": "abc"}' + + +@pytest.mark.asyncio +async def test_quart_ws_connect_default_scheme(quart_app): + async with TestClient(quart_app, timeout=0.1) as client: + async with client.websocket_connect("/ws") as ws: + await ws.send_text("url") + msg = await ws.receive_text() + assert msg.startswith("ws://") + + +@pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") +async def test_quart_ws_connect_custom_scheme(quart_app): + async with TestClient(quart_app, timeout=0.1) as client: + async with client.websocket_connect("/ws", scheme="wss") as ws: + await ws.send_text("url") + msg = await ws.receive_text() + assert msg.startswith("wss://") + + +@pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") +async def test_quart_ws_endpoint_with_immediate_rejection(quart_app): + async with TestClient(quart_app, timeout=0.1) as client: + try: + async with client.websocket_connect("/ws-reject") as ws: + pass + except Exception as e: + thrown_exception = e + + assert ast.literal_eval(str(thrown_exception)) == { + "type": "websocket.close", + "code": starlette.status.WS_1003_UNSUPPORTED_DATA + } + + +@pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") +async def test_quart_invalid_ws_endpoint(quart_app): + async with TestClient(quart_app, timeout=0.1) as client: + try: + async with client.websocket_connect("/invalid") as ws: + pass + except Exception as e: + thrown_exception = e + + assert ast.literal_eval(str(thrown_exception)) == { + "type": "websocket.close", + "code": starlette.status.WS_1000_NORMAL_CLOSURE + } + + @pytest.mark.asyncio async def test_request_stream(starlette_app): from starlette.responses import StreamingResponse From 46e139acc9922c3cf8aff48d5c922e33bd393b44 Mon Sep 17 00:00:00 2001 From: Jake Holland Date: Sun, 12 Jun 2022 21:26:20 +0100 Subject: [PATCH 3/4] Fixed build --- async_asgi_testclient/tests/test_testing.py | 93 +++++++++++++++++---- async_asgi_testclient/websocket.py | 5 +- test-requirements.txt | 2 +- 3 files changed, 79 insertions(+), 21 deletions(-) diff --git a/async_asgi_testclient/tests/test_testing.py b/async_asgi_testclient/tests/test_testing.py index ed92f14..c8cd1cf 100644 --- a/async_asgi_testclient/tests/test_testing.py +++ b/async_asgi_testclient/tests/test_testing.py @@ -1,21 +1,20 @@ -import ast - -import starlette.status -from quart import websocket - from async_asgi_testclient import TestClient from http.cookies import SimpleCookie from json import dumps +from starlette.responses import RedirectResponse +from starlette.responses import StreamingResponse from sys import version_info as PY_VER # noqa +import ast import asyncio import io import pytest +import starlette.status @pytest.fixture def quart_app(): - from quart import Quart, jsonify, request, redirect, Response + from quart import Quart, jsonify, request, redirect, Response, websocket app = Quart(__name__) @@ -97,7 +96,9 @@ async def websocket_endpoint(): @app.websocket("/ws-reject") async def websocket_reject(): - await websocket.close(code=starlette.status.WS_1003_UNSUPPORTED_DATA, reason="some reason") + await websocket.close( + code=starlette.status.WS_1003_UNSUPPORTED_DATA, reason="some reason" + ) yield app @@ -211,6 +212,10 @@ async def echoheaders(request): async def test_query(request): return Response(str(request.query_params)) + @app.route("/redir") + async def redir(request): + return RedirectResponse(request.query_params["path"], status_code=302) + yield app @@ -376,6 +381,29 @@ async def test_set_cookie_in_request(quart_app): assert resp.text == "my-cookie=1234; my-cookie-2=5678" +@pytest.mark.asyncio +async def test_set_cookie_in_request_starlette(starlette_app): + async with TestClient(starlette_app) as client: + resp = await client.post("/set_cookies") + assert resp.status_code == 200 + assert resp.cookies.get_dict() == {"my-cookie": "1234", "my-cookie-2": "5678"} + + # Uses 'custom_cookie_jar' instead of 'client.cookie_jar' + custom_cookie_jar = {"my-cookie": "6666"} + resp = await client.get("/cookies", cookies=custom_cookie_jar) + assert resp.status_code == 200 + assert resp.json() == custom_cookie_jar + + # Uses 'client.cookie_jar' again + resp = await client.get("/cookies") + assert resp.status_code == 200 + assert resp.json() == {"my-cookie": "1234", "my-cookie-2": "5678"} + + resp = await client.get("/cookies-raw") + assert resp.status_code == 200 + assert resp.text == "my-cookie=1234; my-cookie-2=5678" + + @pytest.mark.asyncio @pytest.mark.skipif("PY_VER < (3,7)") async def test_disable_cookies_in_client(quart_app): @@ -478,14 +506,14 @@ async def test_ws_connect_custom_scheme(starlette_app): async def test_ws_endpoint_with_immediate_rejection(starlette_app): async with TestClient(starlette_app, timeout=0.1) as client: try: - async with client.websocket_connect("/ws-reject") as ws: + async with client.websocket_connect("/ws-reject"): pass except Exception as e: thrown_exception = e assert ast.literal_eval(str(thrown_exception)) == { "type": "websocket.close", - "code": starlette.status.WS_1003_UNSUPPORTED_DATA + "code": starlette.status.WS_1003_UNSUPPORTED_DATA, } @@ -493,14 +521,14 @@ async def test_ws_endpoint_with_immediate_rejection(starlette_app): async def test_invalid_ws_endpoint(starlette_app): async with TestClient(starlette_app, timeout=0.1) as client: try: - async with client.websocket_connect("/invalid") as ws: + async with client.websocket_connect("/invalid"): pass except Exception as e: thrown_exception = e assert ast.literal_eval(str(thrown_exception)) == { "type": "websocket.close", - "code": starlette.status.WS_1000_NORMAL_CLOSURE + "code": starlette.status.WS_1000_NORMAL_CLOSURE, } @@ -537,6 +565,7 @@ async def test_quart_ws_connect_inherits_test_client_cookies(quart_app): @pytest.mark.asyncio +@pytest.mark.skipif("PY_VER < (3,7)") async def test_quart_ws_connect_default_scheme(quart_app): async with TestClient(quart_app, timeout=0.1) as client: async with client.websocket_connect("/ws") as ws: @@ -560,14 +589,14 @@ async def test_quart_ws_connect_custom_scheme(quart_app): async def test_quart_ws_endpoint_with_immediate_rejection(quart_app): async with TestClient(quart_app, timeout=0.1) as client: try: - async with client.websocket_connect("/ws-reject") as ws: + async with client.websocket_connect("/ws-reject"): pass except Exception as e: thrown_exception = e assert ast.literal_eval(str(thrown_exception)) == { "type": "websocket.close", - "code": starlette.status.WS_1003_UNSUPPORTED_DATA + "code": starlette.status.WS_1003_UNSUPPORTED_DATA, } @@ -576,14 +605,14 @@ async def test_quart_ws_endpoint_with_immediate_rejection(quart_app): async def test_quart_invalid_ws_endpoint(quart_app): async with TestClient(quart_app, timeout=0.1) as client: try: - async with client.websocket_connect("/invalid") as ws: + async with client.websocket_connect("/invalid"): pass except Exception as e: thrown_exception = e assert ast.literal_eval(str(thrown_exception)) == { "type": "websocket.close", - "code": starlette.status.WS_1000_NORMAL_CLOSURE + "code": starlette.status.WS_1000_NORMAL_CLOSURE, } @@ -664,7 +693,24 @@ async def async_generator(): @pytest.mark.asyncio -async def test_response_stream_crashes(starlette_app): +async def test_response_stream_starlette(starlette_app): + @starlette_app.route("/download_stream") + async def down_stream(_): + async def async_generator(): + chunk = b"X" * 1024 + for _ in range(3): + yield chunk + + return StreamingResponse(async_generator()) + + async with TestClient(starlette_app) as client: + resp = await client.get("/download_stream", stream=False) + assert resp.status_code == 200 + assert len(resp.content) == 3 * 1024 + + +@pytest.mark.asyncio +async def test_response_stream_crashes_starlette(starlette_app): from starlette.responses import StreamingResponse @starlette_app.route("/download_stream_crashes") @@ -701,3 +747,18 @@ async def test_no_follow_redirects(quart_app): async with TestClient(quart_app) as client: resp = await client.get("/redir?path=/", allow_redirects=False) assert resp.status_code == 302 + + +@pytest.mark.asyncio +async def test_follow_redirects_starlette(starlette_app): + async with TestClient(starlette_app) as client: + resp = await client.get("/redir?path=/") + assert resp.status_code == 200 + assert resp.text == "full response" + + +@pytest.mark.asyncio +async def test_no_follow_redirects_starlette(starlette_app): + async with TestClient(starlette_app) as client: + resp = await client.get("/redir?path=/", allow_redirects=False) + assert resp.status_code == 302 diff --git a/async_asgi_testclient/websocket.py b/async_asgi_testclient/websocket.py index e895cd2..84b2d04 100644 --- a/async_asgi_testclient/websocket.py +++ b/async_asgi_testclient/websocket.py @@ -87,10 +87,7 @@ async def receive_json(self): return json.loads(data) async def _receive(self): - try: - return await receive(self.output_queue) - except: - pass + return await receive(self.output_queue) def __aiter__(self): return self diff --git a/test-requirements.txt b/test-requirements.txt index 749a341..3ba301c 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1,4 +1,4 @@ -quart==0.10.0; python_version >= '3.7' +quart==0.17.0; python_version >= '3.7' starlette==0.12.13 python-multipart==0.0.5 pytest==6.2.5 From c3bf21ac971d77eaed41fd155a874813a4069eda Mon Sep 17 00:00:00 2001 From: Jake Holland Date: Mon, 13 Jun 2022 08:27:47 +0100 Subject: [PATCH 4/4] bumped black version --- test-requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test-requirements.txt b/test-requirements.txt index 3ba301c..8816c81 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -4,7 +4,7 @@ python-multipart==0.0.5 pytest==6.2.5 pytest-asyncio==0.15.0 pytest-cov==2.8.1 -black==19.10b0 +black==22.3.0 flake8~=3.8.0 mypy==0.761 isort==4.3.21