From 1b0153d5b8bd0e3bbbf9dceba78be058bd8a220a Mon Sep 17 00:00:00 2001
From: sondale-git <61547150+sondale-git@users.noreply.github.com>
Date: Sat, 5 Feb 2022 21:13:45 +0100
Subject: [PATCH 1/7] gql/transports/requests.py add response header attribute

---
 gql/transport/requests.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/gql/transport/requests.py b/gql/transport/requests.py
index 32e57478..98cdf2dd 100644
--- a/gql/transport/requests.py
+++ b/gql/transport/requests.py
@@ -81,6 +81,8 @@ def __init__(
 
         self.session = None
 
+        self.response_header = None
+
     def connect(self):
 
         if self.session is None:
@@ -217,6 +219,7 @@ def execute(  # type: ignore
         response = self.session.request(
             self.method, self.url, **post_args  # type: ignore
         )
+        self.response_header = response.headers
 
         def raise_response_error(resp: requests.Response, reason: str):
             # We raise a TransportServerError if the status code is 400 or higher

From c3134f478638575f7a55d072600e6e438a2a15e9 Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Thu, 10 Mar 2022 12:34:23 +0100
Subject: [PATCH 2/7] Rename response_header to response_headers

---
 gql/transport/requests.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/gql/transport/requests.py b/gql/transport/requests.py
index 98cdf2dd..a34e7542 100644
--- a/gql/transport/requests.py
+++ b/gql/transport/requests.py
@@ -81,7 +81,7 @@ def __init__(
 
         self.session = None
 
-        self.response_header = None
+        self.response_headers = None
 
     def connect(self):
 
@@ -219,7 +219,7 @@ def execute(  # type: ignore
         response = self.session.request(
             self.method, self.url, **post_args  # type: ignore
         )
-        self.response_header = response.headers
+        self.response_headers = response.headers
 
         def raise_response_error(resp: requests.Response, reason: str):
             # We raise a TransportServerError if the status code is 400 or higher

From 1adaa4038c3561056168ecb765e3a074813670a3 Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Thu, 10 Mar 2022 16:44:55 +0100
Subject: [PATCH 3/7] Adding test for saving headers in transport

---
 tests/test_requests.py | 69 ++++++++++++++++++++++++------------------
 1 file changed, 40 insertions(+), 29 deletions(-)

diff --git a/tests/test_requests.py b/tests/test_requests.py
index 1ed4ca56..7cd7f712 100644
--- a/tests/test_requests.py
+++ b/tests/test_requests.py
@@ -1,3 +1,5 @@
+from typing import Mapping
+
 import pytest
 
 from gql import Client, gql
@@ -38,7 +40,11 @@ async def test_requests_query(event_loop, aiohttp_server, run_sync_test):
     from gql.transport.requests import RequestsHTTPTransport
 
     async def handler(request):
-        return web.Response(text=query1_server_answer, content_type="application/json")
+        return web.Response(
+            text=query1_server_answer,
+            content_type="application/json",
+            headers={"dummy": "test1234"},
+        )
 
     app = web.Application()
     app.router.add_route("POST", "/", handler)
@@ -47,9 +53,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -62,6 +68,11 @@ def test_code():
 
             assert africa["code"] == "AF"
 
+            # Checking response headers are saved in the transport
+            assert hasattr(transport, "response_headers")
+            assert isinstance(transport.response_headers, Mapping)
+            assert transport.response_headers["dummy"] == "test1234"
+
     await run_sync_test(event_loop, server, test_code)
 
 
@@ -84,9 +95,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url, cookies={"cookie1": "val1"})
+        transport = RequestsHTTPTransport(url=url, cookies={"cookie1": "val1"})
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -123,9 +134,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -154,9 +165,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -187,9 +198,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -225,9 +236,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -253,9 +264,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             with pytest.raises(TransportAlreadyConnected):
                 session.transport.connect()
@@ -281,12 +292,12 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
         query = gql(query1_str)
 
         with pytest.raises(TransportClosed):
-            sample_transport.execute(query)
+            transport.execute(query)
 
     await run_sync_test(event_loop, server, test_code)
 
@@ -322,9 +333,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
-        with Client(transport=sample_transport,) as session:
+        with Client(transport=transport) as session:
 
             query = gql(query1_str)
 
@@ -399,10 +410,10 @@ async def single_upload_handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
         with TemporaryFile(file_1_content) as test_file:
-            with Client(transport=sample_transport) as session:
+            with Client(transport=transport) as session:
                 query = gql(file_upload_mutation_1)
 
                 file_path = test_file.filename
@@ -463,10 +474,10 @@ async def single_upload_handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url, headers={"X-Auth": "foobar"})
+        transport = RequestsHTTPTransport(url=url, headers={"X-Auth": "foobar"})
 
         with TemporaryFile(file_1_content) as test_file:
-            with Client(transport=sample_transport) as session:
+            with Client(transport=transport) as session:
                 query = gql(file_upload_mutation_1)
 
                 file_path = test_file.filename
@@ -526,11 +537,11 @@ async def binary_upload_handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = RequestsHTTPTransport(url=url)
+    transport = RequestsHTTPTransport(url=url)
 
     def test_code():
         with TemporaryFile(binary_file_content) as test_file:
-            with Client(transport=sample_transport,) as session:
+            with Client(transport=transport) as session:
 
                 query = gql(file_upload_mutation_1)
 
@@ -617,12 +628,12 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
         with TemporaryFile(file_1_content) as test_file_1:
             with TemporaryFile(file_2_content) as test_file_2:
 
-                with Client(transport=sample_transport,) as session:
+                with Client(transport=transport) as session:
 
                     query = gql(file_upload_mutation_2)
 
@@ -718,11 +729,11 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = RequestsHTTPTransport(url=url)
+        transport = RequestsHTTPTransport(url=url)
 
         with TemporaryFile(file_1_content) as test_file_1:
             with TemporaryFile(file_2_content) as test_file_2:
-                with Client(transport=sample_transport,) as session:
+                with Client(transport=transport) as session:
 
                     query = gql(file_upload_mutation_3)
 

From 24d9444972738cf0487d4d95c6158de2540c704f Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Thu, 10 Mar 2022 16:59:36 +0100
Subject: [PATCH 4/7] Saving response_headers also for AIOHTTPTransport

---
 gql/transport/aiohttp.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/gql/transport/aiohttp.py b/gql/transport/aiohttp.py
index 12c57068..7df60417 100644
--- a/gql/transport/aiohttp.py
+++ b/gql/transport/aiohttp.py
@@ -12,6 +12,7 @@
 from aiohttp.helpers import BasicAuth
 from aiohttp.typedefs import LooseCookies, LooseHeaders
 from graphql import DocumentNode, ExecutionResult, print_ast
+from multidict import CIMultiDictProxy
 
 from ..utils import extract_files
 from .appsync_auth import AppSyncAuthentication
@@ -75,6 +76,7 @@ def __init__(
         self.ssl_close_timeout: Optional[Union[int, float]] = ssl_close_timeout
         self.client_session_args = client_session_args
         self.session: Optional[aiohttp.ClientSession] = None
+        self.response_headers: Optional[CIMultiDictProxy[str]]
 
     async def connect(self) -> None:
         """Coroutine which will create an aiohttp ClientSession() as self.session.
@@ -311,6 +313,9 @@ async def raise_response_error(resp: aiohttp.ClientResponse, reason: str):
             if "errors" not in result and "data" not in result:
                 await raise_response_error(resp, 'No "data" or "errors" keys in answer')
 
+            # Saving latest response headers in the transport
+            self.response_headers = resp.headers
+
             return ExecutionResult(
                 errors=result.get("errors"),
                 data=result.get("data"),

From 544a7d568f31031c60ec07d1a81ffab3d0165ad0 Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Thu, 10 Mar 2022 17:02:57 +0100
Subject: [PATCH 5/7] Adding test for saving headers in AIOHTTPTransport

---
 tests/test_aiohttp.py | 108 +++++++++++++++++++++++-------------------
 1 file changed, 59 insertions(+), 49 deletions(-)

diff --git a/tests/test_aiohttp.py b/tests/test_aiohttp.py
index 030c9134..ab02e8f5 100644
--- a/tests/test_aiohttp.py
+++ b/tests/test_aiohttp.py
@@ -1,5 +1,6 @@
 import io
 import json
+from typing import Mapping
 
 import pytest
 
@@ -45,7 +46,11 @@ async def test_aiohttp_query(event_loop, aiohttp_server):
     from gql.transport.aiohttp import AIOHTTPTransport
 
     async def handler(request):
-        return web.Response(text=query1_server_answer, content_type="application/json")
+        return web.Response(
+            text=query1_server_answer,
+            content_type="application/json",
+            headers={"dummy": "test1234"},
+        )
 
     app = web.Application()
     app.router.add_route("POST", "/", handler)
@@ -53,9 +58,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -68,6 +73,11 @@ async def handler(request):
 
         assert africa["code"] == "AF"
 
+        # Checking response headers are saved in the transport
+        assert hasattr(transport, "response_headers")
+        assert isinstance(transport.response_headers, Mapping)
+        assert transport.response_headers["dummy"] == "test1234"
+
 
 @pytest.mark.asyncio
 async def test_aiohttp_ignore_backend_content_type(event_loop, aiohttp_server):
@@ -83,9 +93,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -115,9 +125,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, cookies={"cookie1": "val1"})
+    transport = AIOHTTPTransport(url=url, cookies={"cookie1": "val1"})
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -150,9 +160,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url)
+    transport = AIOHTTPTransport(url=url)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -177,9 +187,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url)
+    transport = AIOHTTPTransport(url=url)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -208,9 +218,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url)
+    transport = AIOHTTPTransport(url=url)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -259,9 +269,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url)
+    transport = AIOHTTPTransport(url=url)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -285,9 +295,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url)
+    transport = AIOHTTPTransport(url=url)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -310,9 +320,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         with pytest.raises(TransportAlreadyConnected):
             await session.transport.connect()
@@ -332,12 +342,12 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
     query = gql(query1_str)
 
     with pytest.raises(TransportClosed):
-        await sample_transport.execute(query)
+        await transport.execute(query)
 
 
 @pytest.mark.asyncio
@@ -358,11 +368,11 @@ async def handler(request):
     from aiohttp import DummyCookieJar
 
     jar = DummyCookieJar()
-    sample_transport = AIOHTTPTransport(
+    transport = AIOHTTPTransport(
         url=url, timeout=10, client_session_args={"version": "1.1", "cookie_jar": jar}
     )
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -401,9 +411,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         params = {"code": "EU"}
 
@@ -437,9 +447,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         params = {"code": "EU"}
 
@@ -470,9 +480,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = AIOHTTPTransport(url=url)
+        transport = AIOHTTPTransport(url=url)
 
-        client = Client(transport=sample_transport)
+        client = Client(transport=transport)
 
         query = gql(query1_str)
 
@@ -498,9 +508,9 @@ async def handler(request):
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = AIOHTTPTransport(url=url)
+        transport = AIOHTTPTransport(url=url)
 
-        client = Client(transport=sample_transport)
+        client = Client(transport=transport)
 
         query = gql(query1_str)
 
@@ -580,11 +590,11 @@ async def test_aiohttp_file_upload(event_loop, aiohttp_server):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
     with TemporaryFile(file_1_content) as test_file:
 
-        async with Client(transport=sample_transport,) as session:
+        async with Client(transport=transport) as session:
 
             query = gql(file_upload_mutation_1)
 
@@ -618,11 +628,11 @@ async def test_aiohttp_file_upload_without_session(
     url = server.make_url("/")
 
     def test_code():
-        sample_transport = AIOHTTPTransport(url=url, timeout=10)
+        transport = AIOHTTPTransport(url=url, timeout=10)
 
         with TemporaryFile(file_1_content) as test_file:
 
-            client = Client(transport=sample_transport,)
+            client = Client(transport=transport)
 
             query = gql(file_upload_mutation_1)
 
@@ -685,11 +695,11 @@ async def test_aiohttp_binary_file_upload(event_loop, aiohttp_server):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
     with TemporaryFile(binary_file_content) as test_file:
 
-        async with Client(transport=sample_transport,) as session:
+        async with Client(transport=transport) as session:
 
             query = gql(file_upload_mutation_1)
 
@@ -728,9 +738,9 @@ async def binary_data_handler(request):
     url = server.make_url("/")
     binary_data_url = server.make_url("/binary_data")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
         query = gql(file_upload_mutation_1)
         async with ClientSession() as client:
             async with client.get(binary_data_url) as resp:
@@ -758,11 +768,11 @@ async def test_aiohttp_async_generator_upload(event_loop, aiohttp_server):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
     with TemporaryFile(binary_file_content) as test_file:
 
-        async with Client(transport=sample_transport,) as session:
+        async with Client(transport=transport) as session:
 
             query = gql(file_upload_mutation_1)
 
@@ -851,12 +861,12 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
     with TemporaryFile(file_1_content) as test_file_1:
         with TemporaryFile(file_2_content) as test_file_2:
 
-            async with Client(transport=sample_transport,) as session:
+            async with Client(transport=transport) as session:
 
                 query = gql(file_upload_mutation_2)
 
@@ -941,12 +951,12 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
     with TemporaryFile(file_1_content) as test_file_1:
         with TemporaryFile(file_2_content) as test_file_2:
 
-            async with Client(transport=sample_transport,) as session:
+            async with Client(transport=transport) as session:
 
                 query = gql(file_upload_mutation_3)
 
@@ -1098,9 +1108,9 @@ async def handler(request):
 
     url = server.make_url("/")
 
-    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+    transport = AIOHTTPTransport(url=url, timeout=10)
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 
@@ -1126,11 +1136,11 @@ async def handler(request):
 
     assert str(url).startswith("https://")
 
-    sample_transport = AIOHTTPTransport(
+    transport = AIOHTTPTransport(
         url=url, timeout=10, ssl_close_timeout=ssl_close_timeout
     )
 
-    async with Client(transport=sample_transport,) as session:
+    async with Client(transport=transport) as session:
 
         query = gql(query1_str)
 

From 8b0d685cba93393348c3b2d7dd6e516eedd8d50c Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Thu, 10 Mar 2022 17:21:04 +0100
Subject: [PATCH 6/7] Saving response_headers in WebsocketsTransport + test

---
 gql/transport/websockets_base.py |  6 +++-
 tests/conftest.py                |  3 ++
 tests/test_websocket_query.py    | 61 +++++++++++++++++---------------
 3 files changed, 41 insertions(+), 29 deletions(-)

diff --git a/gql/transport/websockets_base.py b/gql/transport/websockets_base.py
index 151e444e..7a83c47f 100644
--- a/gql/transport/websockets_base.py
+++ b/gql/transport/websockets_base.py
@@ -9,7 +9,7 @@
 import websockets
 from graphql import DocumentNode, ExecutionResult
 from websockets.client import WebSocketClientProtocol
-from websockets.datastructures import HeadersLike
+from websockets.datastructures import Headers, HeadersLike
 from websockets.exceptions import ConnectionClosed
 from websockets.typing import Data, Subprotocol
 
@@ -169,6 +169,8 @@ def __init__(
         # The list of supported subprotocols should be defined in the subclass
         self.supported_subprotocols: List[Subprotocol] = []
 
+        self.response_headers: Optional[Headers] = None
+
     async def _initialize(self):
         """Hook to send the initialization messages after the connection
         and potentially wait for the backend ack.
@@ -495,6 +497,8 @@ async def connect(self) -> None:
 
             self.websocket = cast(WebSocketClientProtocol, self.websocket)
 
+            self.response_headers = self.websocket.response_headers
+
             # Run the after_connect hook of the subclass
             await self._after_connect()
 
diff --git a/tests/conftest.py b/tests/conftest.py
index fbb881ec..8d0b95ba 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -174,6 +174,9 @@ async def start(self, handler, extra_serve_args=None):
             self.testcert, ssl_context = get_localhost_ssl_context()
             extra_serve_args["ssl"] = ssl_context
 
+        # Adding dummy response headers
+        extra_serve_args["extra_headers"] = {"dummy": "test1234"}
+
         # Start a server with a random open port
         self.start_server = websockets.server.serve(
             handler, "127.0.0.1", 0, **extra_serve_args
diff --git a/tests/test_websocket_query.py b/tests/test_websocket_query.py
index 4e51f161..2382f157 100644
--- a/tests/test_websocket_query.py
+++ b/tests/test_websocket_query.py
@@ -2,7 +2,7 @@
 import json
 import ssl
 import sys
-from typing import Dict
+from typing import Dict, Mapping
 
 import pytest
 
@@ -58,12 +58,12 @@ async def test_websocket_starting_client_in_context_manager(event_loop, server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url)
+    transport = WebsocketsTransport(url=url)
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
 
         assert isinstance(
-            sample_transport.websocket, websockets.client.WebSocketClientProtocol
+            transport.websocket, websockets.client.WebSocketClientProtocol
         )
 
         query1 = gql(query1_str)
@@ -80,8 +80,13 @@ async def test_websocket_starting_client_in_context_manager(event_loop, server):
 
         assert africa["code"] == "AF"
 
+        # Checking response headers are saved in the transport
+        assert hasattr(transport, "response_headers")
+        assert isinstance(transport.response_headers, Mapping)
+        assert transport.response_headers["dummy"] == "test1234"
+
     # Check client is disconnect here
-    assert sample_transport.websocket is None
+    assert transport.websocket is None
 
 
 @pytest.mark.asyncio
@@ -98,12 +103,12 @@ async def test_websocket_using_ssl_connection(event_loop, ws_ssl_server):
     ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
     ssl_context.load_verify_locations(ws_ssl_server.testcert)
 
-    sample_transport = WebsocketsTransport(url=url, ssl=ssl_context)
+    transport = WebsocketsTransport(url=url, ssl=ssl_context)
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
 
         assert isinstance(
-            sample_transport.websocket, websockets.client.WebSocketClientProtocol
+            transport.websocket, websockets.client.WebSocketClientProtocol
         )
 
         query1 = gql(query1_str)
@@ -121,7 +126,7 @@ async def test_websocket_using_ssl_connection(event_loop, ws_ssl_server):
         assert africa["code"] == "AF"
 
     # Check client is disconnect here
-    assert sample_transport.websocket is None
+    assert transport.websocket is None
 
 
 @pytest.mark.asyncio
@@ -301,19 +306,19 @@ async def test_websocket_multiple_connections_in_series(event_loop, server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url)
+    transport = WebsocketsTransport(url=url)
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
         await assert_client_is_working(session)
 
     # Check client is disconnect here
-    assert sample_transport.websocket is None
+    assert transport.websocket is None
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
         await assert_client_is_working(session)
 
     # Check client is disconnect here
-    assert sample_transport.websocket is None
+    assert transport.websocket is None
 
 
 @pytest.mark.asyncio
@@ -325,8 +330,8 @@ async def test_websocket_multiple_connections_in_parallel(event_loop, server):
     print(f"url = {url}")
 
     async def task_coro():
-        sample_transport = WebsocketsTransport(url=url)
-        async with Client(transport=sample_transport) as session:
+        transport = WebsocketsTransport(url=url)
+        async with Client(transport=transport) as session:
             await assert_client_is_working(session)
 
     task1 = asyncio.ensure_future(task_coro())
@@ -345,12 +350,12 @@ async def test_websocket_trying_to_connect_to_already_connected_transport(
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url)
-    async with Client(transport=sample_transport) as session:
+    transport = WebsocketsTransport(url=url)
+    async with Client(transport=transport) as session:
         await assert_client_is_working(session)
 
         with pytest.raises(TransportAlreadyConnected):
-            async with Client(transport=sample_transport):
+            async with Client(transport=transport):
                 pass
 
 
@@ -395,9 +400,9 @@ async def test_websocket_connect_success_with_authentication_in_connection_init(
 
     init_payload = {"Authorization": 12345}
 
-    sample_transport = WebsocketsTransport(url=url, init_payload=init_payload)
+    transport = WebsocketsTransport(url=url, init_payload=init_payload)
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
 
         query1 = gql(query_str)
 
@@ -428,10 +433,10 @@ async def test_websocket_connect_failed_with_authentication_in_connection_init(
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url, init_payload=init_payload)
+    transport = WebsocketsTransport(url=url, init_payload=init_payload)
 
     with pytest.raises(TransportServerError):
-        async with Client(transport=sample_transport) as session:
+        async with Client(transport=transport) as session:
             query1 = gql(query_str)
 
             await session.execute(query1)
@@ -444,9 +449,9 @@ def test_websocket_execute_sync(server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
     print(f"url = {url}")
 
-    sample_transport = WebsocketsTransport(url=url)
+    transport = WebsocketsTransport(url=url)
 
-    client = Client(transport=sample_transport)
+    client = Client(transport=transport)
 
     query1 = gql(query1_str)
 
@@ -476,7 +481,7 @@ def test_websocket_execute_sync(server):
     assert africa["code"] == "AF"
 
     # Check client is disconnect here
-    assert sample_transport.websocket is None
+    assert transport.websocket is None
 
 
 @pytest.mark.asyncio
@@ -487,11 +492,11 @@ async def test_websocket_add_extra_parameters_to_connect(event_loop, server):
     url = f"ws://{server.hostname}:{server.port}/graphql"
 
     # Increase max payload size to avoid websockets.exceptions.PayloadTooBig exceptions
-    sample_transport = WebsocketsTransport(url=url, connect_args={"max_size": 2 ** 21})
+    transport = WebsocketsTransport(url=url, connect_args={"max_size": 2 ** 21})
 
     query = gql(query1_str)
 
-    async with Client(transport=sample_transport) as session:
+    async with Client(transport=transport) as session:
         await session.execute(query)
 
 

From ca098b12655ea106ea9fe2480833efa2f561f429 Mon Sep 17 00:00:00 2001
From: Hanusz Leszek <leszek.hanusz@gmail.com>
Date: Thu, 10 Mar 2022 17:28:07 +0100
Subject: [PATCH 7/7] Add docs

---
 docs/usage/headers.rst | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/usage/headers.rst b/docs/usage/headers.rst
index 23af64a7..b41c8b43 100644
--- a/docs/usage/headers.rst
+++ b/docs/usage/headers.rst
@@ -6,3 +6,5 @@ If you want to add additional http headers for your connection, you can specify
 .. code-block:: python
 
     transport = AIOHTTPTransport(url='YOUR_URL', headers={'Authorization': 'token'})
+
+After the connection, the latest response headers can be found in :code:`transport.response_headers`