From ccec58eccff58db4c977db559ecd46bbbe2c5765 Mon Sep 17 00:00:00 2001 From: Irfanuddin Date: Fri, 28 Oct 2022 20:29:51 +0530 Subject: [PATCH 1/4] add tests for shutdown --- tests/protocols/test_websocket.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index 904875556..e04eaa777 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -5,6 +5,7 @@ from tests.protocols.test_http import HTTP_PROTOCOLS from tests.utils import run_server +from uvicorn import Server from uvicorn.config import Config from uvicorn.protocols.websockets.wsproto_impl import WSProtocol @@ -20,6 +21,7 @@ ONLY_WEBSOCKETPROTOCOL = [p for p in [WebSocketProtocol] if p is not None] +ONLY_WSPROTO = [p for p in [WSProtocol] if p is not None] WS_PROTOCOLS = [p for p in [WSProtocol, WebSocketProtocol] if p is not None] pytestmark = pytest.mark.skipif( websockets is None, reason="This test needs the websockets module" @@ -658,3 +660,32 @@ async def send_text(url): await send_text("ws://127.0.0.1:8000") assert frames == [b"abc", b"abc", b"abc"] + + +@pytest.mark.anyio +@pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS) +@pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS) +async def test_server_shutdown_when_connection_active_in_ws( + ws_protocol_cls, http_protocol_cls +): + class App(WebSocketResponse): + async def websocket_connect(self, message): + await self.send({"type": "websocket.accept"}) + + config = Config( + app=App, + ws=ws_protocol_cls, + http=http_protocol_cls, + lifespan="off", + ) + server = Server(config=config) + cancel_handle = asyncio.ensure_future(server.serve(sockets=None)) + await asyncio.sleep(0.1) + async with websockets.connect("ws://127.0.0.1:8000") as websocket: + ws_conn = list(server.server_state.connections)[0] + ws_conn.shutdown() + await asyncio.sleep(0.1) + assert websocket.close_code == 1012 + assert ws_conn.transport.is_closing() + await server.shutdown() + cancel_handle.cancel() From ff692012518e348c6578f443db0e221c79338628 Mon Sep 17 00:00:00 2001 From: Irfanuddin Date: Fri, 28 Oct 2022 22:48:30 +0530 Subject: [PATCH 2/4] add proper connection close in websockets --- tests/protocols/test_websocket.py | 4 ++-- uvicorn/protocols/websockets/websockets_impl.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index e04eaa777..4fdb37102 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -679,7 +679,7 @@ async def websocket_connect(self, message): lifespan="off", ) server = Server(config=config) - cancel_handle = asyncio.ensure_future(server.serve(sockets=None)) + task = asyncio.create_task(server.serve(sockets=None)) await asyncio.sleep(0.1) async with websockets.connect("ws://127.0.0.1:8000") as websocket: ws_conn = list(server.server_state.connections)[0] @@ -688,4 +688,4 @@ async def websocket_connect(self, message): assert websocket.close_code == 1012 assert ws_conn.transport.is_closing() await server.shutdown() - cancel_handle.cancel() + task.cancel() diff --git a/uvicorn/protocols/websockets/websockets_impl.py b/uvicorn/protocols/websockets/websockets_impl.py index b77485136..f19938d53 100644 --- a/uvicorn/protocols/websockets/websockets_impl.py +++ b/uvicorn/protocols/websockets/websockets_impl.py @@ -137,7 +137,9 @@ def connection_lost(self, exc: Optional[Exception]) -> None: def shutdown(self) -> None: self.ws_server.closing = True - self.transport.close() + task = asyncio.create_task(self.close(code=1012)) + task.add_done_callback(self.on_task_complete) + self.tasks.add(task) def on_task_complete(self, task: asyncio.Task) -> None: self.tasks.discard(task) From 58333357d7b038f113e9b4bb1499a413ffdaf06e Mon Sep 17 00:00:00 2001 From: Irfanuddin Date: Fri, 28 Oct 2022 23:00:06 +0530 Subject: [PATCH 3/4] generic test name --- tests/protocols/test_websocket.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/protocols/test_websocket.py b/tests/protocols/test_websocket.py index c7d302186..e4c1e39b1 100644 --- a/tests/protocols/test_websocket.py +++ b/tests/protocols/test_websocket.py @@ -762,7 +762,7 @@ async def open_connection(url): @pytest.mark.anyio @pytest.mark.parametrize("ws_protocol_cls", WS_PROTOCOLS) @pytest.mark.parametrize("http_protocol_cls", HTTP_PROTOCOLS) -async def test_server_shutdown_when_connection_active_in_ws( +async def test_server_shutdown_when_connection_active( ws_protocol_cls, http_protocol_cls ): class App(WebSocketResponse): From 4bc06b2120b57dd8913a82539b30ef136180740d Mon Sep 17 00:00:00 2001 From: Irfanuddin Date: Sat, 29 Oct 2022 19:34:31 +0530 Subject: [PATCH 4/4] update fail_under --- setup.cfg | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.cfg b/setup.cfg index 46f4e3b99..99c7a48c3 100644 --- a/setup.cfg +++ b/setup.cfg @@ -82,7 +82,7 @@ plugins = [coverage:report] precision = 2 -fail_under = 97.82 +fail_under = 97.97 show_missing = true skip_covered = true exclude_lines =