Bump websockets to >=14.2

Using new websockets asyncio implementation

Author: leszekhanusz

gql/transport/common/adapters/websockets.py

@@ -3,7 +3,7 @@
 from typing import Any, Dict, Optional, Union
 
 import websockets
-from websockets.client import WebSocketClientProtocol
+from websockets import ClientConnection
 from websockets.datastructures import Headers, HeadersLike
 
 from ...exceptions import TransportConnectionFailed, TransportProtocolError
@@ -40,7 +40,7 @@ def __init__(
         self._headers: Optional[HeadersLike] = headers
         self.ssl = ssl
 
-        self.websocket: Optional[WebSocketClientProtocol] = None
+        self.websocket: Optional[ClientConnection] = None
         self._response_headers: Optional[Headers] = None
 
     async def connect(self) -> None:
@@ -57,7 +57,7 @@ async def connect(self) -> None:
         # Set default arguments used in the websockets.connect call
         connect_args: Dict[str, Any] = {
             "ssl": ssl,
-            "extra_headers": self.headers,
+            "additional_headers": self.headers,
         }
 
         if self.subprotocols:
@@ -68,11 +68,13 @@ async def connect(self) -> None:
 
         # Connection to the specified url
         try:
-            self.websocket = await websockets.client.connect(self.url, **connect_args)
+            self.websocket = await websockets.connect(self.url, **connect_args)
         except Exception as e:
             raise TransportConnectionFailed("Connect failed") from e
 
-        self._response_headers = self.websocket.response_headers
+        assert self.websocket.response is not None
+
+        self._response_headers = self.websocket.response.headers
 
     async def send(self, message: str) -> None:
         """Send message to the WebSocket server.

gql/transport/common/base.py

@@ -482,6 +482,10 @@ async def _close_coro(self, e: Exception, clean_close: bool = True) -> None:
             # We should always have an active websocket connection here
             assert self._connected
 
+            # Saving exception to raise it later if trying to use the transport
+            # after it has already closed.
+            self.close_exception = e
+
             # Properly shut down liveness checker if enabled
             if self.check_keep_alive_task is not None:
                 # More info: https://stackoverflow.com/a/43810272/1113207
@@ -492,18 +496,16 @@ async def _close_coro(self, e: Exception, clean_close: bool = True) -> None:
             # Calling the subclass close hook
             await self._close_hook()
 
-            # Saving exception to raise it later if trying to use the transport
-            # after it has already closed.
-            self.close_exception = e
-
             if clean_close:
                 log.debug("_close_coro: starting clean_close")
                 try:
                     await self._clean_close(e)
                 except Exception as exc:  # pragma: no cover
                     log.warning("Ignoring exception in _clean_close: " + repr(exc))
 
-            log.debug("_close_coro: sending exception to listeners")
+            log.debug(
+                f"_close_coro: sending exception to {len(self.listeners)} listeners"
+            )
 
             # Send an exception to all remaining listeners
             for query_id, listener in self.listeners.items():

setup.py

@@ -51,7 +51,7 @@
 ]
 
 install_websockets_requires = [
-    "websockets>=10.1,<14",
+    "websockets>=14.2,<16",
 ]
 
 install_botocore_requires = [

tests/conftest.py

@@ -197,7 +197,7 @@ def __init__(self, with_ssl: bool = False):
 
     async def start(self, handler, extra_serve_args=None):
 
-        import websockets.server
+        import websockets
 
         print("Starting server")
 
@@ -209,16 +209,21 @@ async def start(self, handler, extra_serve_args=None):
             extra_serve_args["ssl"] = ssl_context
 
         # Adding dummy response headers
-        extra_serve_args["extra_headers"] = {"dummy": "test1234"}
+        extra_headers = {"dummy": "test1234"}
+
+        def process_response(connection, request, response):
+            response.headers.update(extra_headers)
+            return response
 
         # Start a server with a random open port
-        self.start_server = websockets.server.serve(
-            handler, "127.0.0.1", 0, **extra_serve_args
+        self.server = await websockets.serve(
+            handler,
+            "127.0.0.1",
+            0,
+            process_response=process_response,
+            **extra_serve_args,
         )
 
-        # Wait that the server is started
-        self.server = await self.start_server
-
         # Get hostname and port
         hostname, port = self.server.sockets[0].getsockname()[:2]  # type: ignore
         assert hostname == "127.0.0.1"
@@ -603,32 +608,14 @@ async def graphqlws_server(request):
 
     subprotocol = "graphql-transport-ws"
 
-    from websockets.server import WebSocketServerProtocol
-
-    class CustomSubprotocol(WebSocketServerProtocol):
-        def select_subprotocol(self, client_subprotocols, server_subprotocols):
-            print(f"Client subprotocols: {client_subprotocols!r}")
-            print(f"Server subprotocols: {server_subprotocols!r}")
-
-            return subprotocol
-
-        def process_subprotocol(self, headers, available_subprotocols):
-            # Overwriting available subprotocols
-            available_subprotocols = [subprotocol]
-
-            print(f"headers: {headers!r}")
-            # print (f"Available subprotocols: {available_subprotocols!r}")
-
-            return super().process_subprotocol(headers, available_subprotocols)
-
     server_handler = get_server_handler(request)
 
     try:
         test_server = WebSocketServer()
 
         # Starting the server with the fixture param as the handler function
         await test_server.start(
-            server_handler, extra_serve_args={"create_protocol": CustomSubprotocol}
+            server_handler, extra_serve_args={"subprotocols": [subprotocol]}
         )
 
         yield test_server

tests/test_aiohttp_websocket_exceptions.py

@@ -118,10 +118,10 @@ async def test_aiohttp_websocket_server_does_not_send_ack(server, query_str):
 
     url = f"ws://{server.hostname}:{server.port}/graphql"
 
-    sample_transport = AIOHTTPWebsocketsTransport(url=url, ack_timeout=1)
+    transport = AIOHTTPWebsocketsTransport(url=url, ack_timeout=1)
 
     with pytest.raises(asyncio.TimeoutError):
-        async with Client(transport=sample_transport):
+        async with Client(transport=transport):
             pass
 
 
@@ -261,10 +261,10 @@ async def test_aiohttp_websocket_server_does_not_ack(server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = AIOHTTPWebsocketsTransport(url=url)
+    transport = AIOHTTPWebsocketsTransport(url=url)
 
     with pytest.raises(TransportProtocolError):
-        async with Client(transport=sample_transport):
+        async with Client(transport=transport):
             pass
 
 
@@ -281,10 +281,10 @@ async def test_aiohttp_websocket_server_closing_directly(server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = AIOHTTPWebsocketsTransport(url=url)
+    transport = AIOHTTPWebsocketsTransport(url=url)
 
     with pytest.raises(TransportConnectionFailed):
-        async with Client(transport=sample_transport):
+        async with Client(transport=transport):
             pass
 
 
@@ -323,10 +323,10 @@ async def test_aiohttp_websocket_server_sending_invalid_query_errors(server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = AIOHTTPWebsocketsTransport(url=url)
+    transport = AIOHTTPWebsocketsTransport(url=url)
 
     # Invalid server message is ignored
-    async with Client(transport=sample_transport):
+    async with Client(transport=transport):
         await asyncio.sleep(2 * MS)
 
 
@@ -342,9 +342,9 @@ async def test_aiohttp_websocket_non_regression_bug_105(server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = AIOHTTPWebsocketsTransport(url=url)
+    transport = AIOHTTPWebsocketsTransport(url=url)
 
-    client = Client(transport=sample_transport)
+    client = Client(transport=transport)
 
     # Create a coroutine which start the connection with the transport but does nothing
     async def client_connect(client):

tests/test_appsync_websockets.py

@@ -139,7 +139,7 @@ async def realtime_appsync_server_template(ws):
                 )
                 return
 
-            path = ws.path
+            path = ws.request.path
 
             print(f"path = {path}")
 

tests/test_graphqlws_exceptions.py

@@ -111,10 +111,10 @@ async def test_graphqlws_server_does_not_send_ack(graphqlws_server, query_str):
 
     url = f"ws://{graphqlws_server.hostname}:{graphqlws_server.port}/graphql"
 
-    sample_transport = WebsocketsTransport(url=url, ack_timeout=1)
+    transport = WebsocketsTransport(url=url, ack_timeout=1)
 
     with pytest.raises(asyncio.TimeoutError):
-        async with Client(transport=sample_transport):
+        async with Client(transport=transport):
             pass
 
 
@@ -212,10 +212,10 @@ async def test_graphqlws_server_does_not_ack(graphqlws_server):
     url = f"ws://{graphqlws_server.hostname}:{graphqlws_server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url)
+    transport = WebsocketsTransport(url=url)
 
     with pytest.raises(TransportProtocolError):
-        async with Client(transport=sample_transport):
+        async with Client(transport=transport):
             pass
 
 
@@ -231,10 +231,10 @@ async def test_graphqlws_server_closing_directly(graphqlws_server):
     url = f"ws://{graphqlws_server.hostname}:{graphqlws_server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url)
+    transport = WebsocketsTransport(url=url)
 
     with pytest.raises(TransportConnectionFailed):
-        async with Client(transport=sample_transport):
+        async with Client(transport=transport):
             pass
 
 
@@ -251,7 +251,7 @@ async def test_graphqlws_server_closing_after_ack(client_and_graphqlws_server):
 
     query = gql("query { hello }")
 
-    with pytest.raises(TransportConnectionFailed):
+    with pytest.raises(TransportClosed):
         await session.execute(query)
 
     await session.transport.wait_closed()

tests/test_websocket_exceptions.py

@@ -296,7 +296,7 @@ async def test_websocket_server_closing_after_ack(client_and_server):
 
     query = gql("query { hello }")
 
-    with pytest.raises(TransportConnectionFailed):
+    with pytest.raises(TransportClosed):
         await session.execute(query)
 
     await session.transport.wait_closed()

tests/test_websocket_query.py

@@ -112,9 +112,7 @@ async def test_websocket_using_ssl_connection(ws_ssl_server):
 
     async with Client(transport=transport) as session:
 
-        assert isinstance(
-            transport.adapter.websocket, websockets.client.WebSocketClientProtocol
-        )
+        assert isinstance(transport.adapter.websocket, websockets.ClientConnection)
 
         query1 = gql(query1_str)