From 9f932159e084fe2686febced979ca21f4b75b745 Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Sun, 14 Jul 2024 18:27:01 +0200
Subject: [PATCH] Fix properly exiting the WebsocketsTransport when a
 connection_error is received during init (#486)

---
 gql/transport/websockets_base.py | 6 +++++-
 tests/test_websocket_query.py    | 3 +++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/gql/transport/websockets_base.py b/gql/transport/websockets_base.py
index 45c96d3e..5c7713e9 100644
--- a/gql/transport/websockets_base.py
+++ b/gql/transport/websockets_base.py
@@ -512,7 +512,11 @@ async def connect(self) -> None:
                 await self._initialize()
             except ConnectionClosed as e:
                 raise e
-            except (TransportProtocolError, asyncio.TimeoutError) as e:
+            except (
+                TransportProtocolError,
+                TransportServerError,
+                asyncio.TimeoutError,
+            ) as e:
                 await self._fail(e, clean_close=False)
                 raise e
 
diff --git a/tests/test_websocket_query.py b/tests/test_websocket_query.py
index e8b7a022..d2270e7d 100644
--- a/tests/test_websocket_query.py
+++ b/tests/test_websocket_query.py
@@ -441,6 +441,9 @@ async def test_websocket_connect_failed_with_authentication_in_connection_init(
 
             await session.execute(query1)
 
+    await asyncio.sleep(1)
+    assert transport.websocket is None
+
 
 @pytest.mark.parametrize("server", [server1_answers], indirect=True)
 def test_websocket_execute_sync(server):