diff --git a/async_asgi_testclient/tests/test_testing.py b/async_asgi_testclient/tests/test_testing.py index c48a9a0..c8cd1cf 100644 --- a/async_asgi_testclient/tests/test_testing.py +++ b/async_asgi_testclient/tests/test_testing.py @@ -1,16 +1,20 @@ 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__) @@ -80,6 +84,22 @@ 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 @@ -108,6 +128,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") @@ -187,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 @@ -352,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): @@ -450,6 +502,120 @@ 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"): + 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"): + 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 +@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 +@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: + 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"): + 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"): + 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 @@ -527,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") @@ -564,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 4f75024..84b2d04 100644 --- a/async_asgi_testclient/websocket.py +++ b/async_asgi_testclient/websocket.py @@ -131,4 +131,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) diff --git a/test-requirements.txt b/test-requirements.txt index 749a341..8816c81 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -1,10 +1,10 @@ -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 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