🐛 Fix: websockets 驱动器连接关闭 code 获取错误 (#2537)

tests/fake_server.py

@@ -1,9 +1,15 @@
 import json
 import base64
+import socket
 from typing import Dict, List, Union, TypeVar
 
+from wsproto.events import Ping
 from werkzeug import Request, Response
 from werkzeug.datastructures import MultiDict
+from wsproto.frame_protocol import CloseReason
+from wsproto.events import Request as WSRequest
+from wsproto import WSConnection, ConnectionType
+from wsproto.events import TextMessage, BytesMessage, CloseConnection, AcceptConnection
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -29,8 +35,7 @@ def flattern(d: "MultiDict[K, V]") -> Dict[K, Union[V, List[V]]]:
     return {k: v[0] if len(v) == 1 else v for k, v in d.to_dict(flat=False).items()}
 
 
-@Request.application
-def request_handler(request: Request) -> Response:
+def http_echo(request: Request) -> Response:
     try:
         _json = json.loads(request.data.decode("utf-8"))
     except (ValueError, TypeError):
@@ -67,3 +72,65 @@ def request_handler(request: Request) -> Response:
         status=200,
         content_type="application/json",
     )
+
+
+def websocket_echo(request: Request) -> Response:
+    stream = request.environ["werkzeug.socket"]
+
+    ws = WSConnection(ConnectionType.SERVER)
+
+    in_data = b"GET %s HTTP/1.1\r\n" % request.path.encode("utf-8")
+    for header, value in request.headers.items():
+        in_data += f"{header}: {value}\r\n".encode()
+    in_data += b"\r\n"
+
+    ws.receive_data(in_data)
+
+    running: bool = True
+    while True:
+        out_data = b""
+
+        for event in ws.events():
+            if isinstance(event, WSRequest):
+                out_data += ws.send(AcceptConnection())
+            elif isinstance(event, CloseConnection):
+                out_data += ws.send(event.response())
+                running = False
+            elif isinstance(event, Ping):
+                out_data += ws.send(event.response())
+            elif isinstance(event, TextMessage):
+                if event.data == "quit":
+                    out_data += ws.send(
+                        CloseConnection(CloseReason.NORMAL_CLOSURE, "bye")
+                    )
+                    running = False
+                else:
+                    out_data += ws.send(TextMessage(data=event.data))
+            elif isinstance(event, BytesMessage):
+                if event.data == b"quit":
+                    out_data += ws.send(
+                        CloseConnection(CloseReason.NORMAL_CLOSURE, "bye")
+                    )
+                    running = False
+                else:
+                    out_data += ws.send(BytesMessage(data=event.data))
+
+        if out_data:
+            stream.send(out_data)
+
+        if not running:
+            break
+
+        in_data = stream.recv(4096)
+        ws.receive_data(in_data)
+
+    stream.shutdown(socket.SHUT_RDWR)
+    return Response("", status=204)
+
+
+@Request.application
+def request_handler(request: Request) -> Response:
+    if request.headers.get("Connection") == "Upgrade":
+        return websocket_echo(request)
+    else:
+        return http_echo(request)

tests/test_driver.py

@@ -131,7 +131,7 @@ async def test_websocket_server(app: App, driver: Driver):
         assert data == b"ping"
         await ws.send(b"pong")
 
-        with pytest.raises(WebSocketClosed):
+        with pytest.raises(WebSocketClosed, match=r"code=1000"):
             await ws.receive()
 
     ws_setup = WebSocketServerSetup(URL("/ws_test"), "ws_test", _handle_ws)
@@ -152,7 +152,7 @@ async def test_websocket_server(app: App, driver: Driver):
             await ws.send_bytes(b"ping")
             assert await ws.receive_bytes() == b"pong"
 
-            await ws.close()
+            await ws.close(code=1000)
 
     await asyncio.sleep(1)
 
@@ -315,9 +315,29 @@ async def test_http_client(driver: Driver, server_url: URL):
     ],
     indirect=True,
 )
-async def test_websocket_client(driver: Driver):
+async def test_websocket_client(driver: Driver, server_url: URL):
     assert isinstance(driver, WebSocketClientMixin)
 
+    request = Request("GET", server_url.with_scheme("ws"))
+    async with driver.websocket(request) as ws:
+        await ws.send("test")
+        assert await ws.receive() == "test"
+
+        await ws.send(b"test")
+        assert await ws.receive() == b"test"
+
+        await ws.send_text("test")
+        assert await ws.receive_text() == "test"
+
+        await ws.send_bytes(b"test")
+        assert await ws.receive_bytes() == b"test"
+
+        await ws.send("quit")
+        with pytest.raises(WebSocketClosed, match=r"code=1000"):
+            await ws.receive()
+
+    await asyncio.sleep(1)
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(