Merge branch 'master' into chore_update_graphql_core_to_3_3_0

leszekhanusz · web-flow · commit 71c5c1a95a71 · 2022-11-07T22:26:29.000+01:00
diff --git a/gql/transport/aiohttp.py b/gql/transport/aiohttp.py
@@ -295,6 +295,9 @@ async def execute(
 
         async with self.session.post(self.url, ssl=self.ssl, **post_args) as resp:
 
+            # Saving latest response headers in the transport
+            self.response_headers = resp.headers
+
             async def raise_response_error(resp: aiohttp.ClientResponse, reason: str):
                 # We raise a TransportServerError if the status code is 400 or higher
                 # We raise a TransportProtocolError in the other cases
@@ -325,9 +328,6 @@ 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"),
diff --git a/tests/test_aiohttp.py b/tests/test_aiohttp.py
@@ -172,6 +172,53 @@ async def handler(request):
         assert "401, message='Unauthorized'" in str(exc_info.value)
 
 
+@pytest.mark.asyncio
+async def test_aiohttp_error_code_429(event_loop, aiohttp_server):
+    from aiohttp import web
+    from gql.transport.aiohttp import AIOHTTPTransport
+
+    async def handler(request):
+        # Will generate http error code 429
+        return web.Response(
+            text="""
+<html>
+  <head>
+     <title>Too Many Requests</title>
+  </head>
+  <body>
+     <h1>Too Many Requests</h1>
+     <p>I only allow 50 requests per hour to this Web site per
+        logged in user.  Try again soon.</p>
+  </body>
+</html>""",
+            content_type="text/html",
+            status=429,
+            headers={"Retry-After": "3600"},
+        )
+
+    app = web.Application()
+    app.router.add_route("POST", "/", handler)
+    server = await aiohttp_server(app)
+
+    url = server.make_url("/")
+
+    transport = AIOHTTPTransport(url=url)
+
+    async with Client(transport=transport) as session:
+
+        query = gql(query1_str)
+
+        with pytest.raises(TransportServerError) as exc_info:
+            await session.execute(query)
+
+        assert "429, message='Too Many Requests'" in str(exc_info.value)
+
+        # Checking response headers are saved in the transport
+        assert hasattr(transport, "response_headers")
+        assert isinstance(transport.response_headers, Mapping)
+        assert transport.response_headers["Retry-After"] == "3600"
+
+
 @pytest.mark.asyncio
 async def test_aiohttp_error_code_500(event_loop, aiohttp_server):
     from aiohttp import web
diff --git a/tests/test_requests.py b/tests/test_requests.py
@@ -148,6 +148,55 @@ def test_code():
     await run_sync_test(event_loop, server, test_code)
 
 
+@pytest.mark.aiohttp
+@pytest.mark.asyncio
+async def test_requests_error_code_429(event_loop, aiohttp_server, run_sync_test):
+    from aiohttp import web
+    from gql.transport.requests import RequestsHTTPTransport
+
+    async def handler(request):
+        # Will generate http error code 429
+        return web.Response(
+            text="""
+<html>
+  <head>
+     <title>Too Many Requests</title>
+  </head>
+  <body>
+     <h1>Too Many Requests</h1>
+     <p>I only allow 50 requests per hour to this Web site per
+        logged in user.  Try again soon.</p>
+  </body>
+</html>""",
+            content_type="text/html",
+            status=429,
+            headers={"Retry-After": "3600"},
+        )
+
+    app = web.Application()
+    app.router.add_route("POST", "/", handler)
+    server = await aiohttp_server(app)
+
+    url = server.make_url("/")
+
+    def test_code():
+        transport = RequestsHTTPTransport(url=url)
+
+        with Client(transport=transport) as session:
+
+            query = gql(query1_str)
+
+            with pytest.raises(TransportServerError) as exc_info:
+                session.execute(query)
+
+            assert "429, message='Too Many Requests'" in str(exc_info.value)
+
+        # Checking response headers are saved in the transport
+        assert hasattr(transport, "response_headers")
+        assert isinstance(transport.response_headers, Mapping)
+        assert transport.response_headers["Retry-After"] == "3600"
+
+
 @pytest.mark.aiohttp
 @pytest.mark.asyncio
 async def test_requests_error_code_500(event_loop, aiohttp_server, run_sync_test):
