Fix PyPy tests

Author: leszekhanusz

gql/transport/common/base.py

@@ -317,6 +317,8 @@ async def subscribe(
             if listener.send_stop:
                 await self._stop_listener(query_id)
                 listener.send_stop = False
+            if isinstance(e, GeneratorExit):
+                raise e
 
         finally:
             log.debug(f"In subscribe finally for query_id {query_id}")
@@ -345,6 +347,11 @@ async def execute(
             first_result = result
             break
 
+        # Apparently, on pypy the GeneratorExit exception is not raised after a break
+        # --> the clean_close has to time out
+        # We still need to manually close the async generator
+        await generator.aclose()
+
         if first_result is None:
             raise TransportQueryError(
                 "Query completed without any answer received from the server"
@@ -445,7 +452,6 @@ async def _clean_close(self, e: Exception) -> None:
 
         # Send 'stop' message for all current queries
         for query_id, listener in self.listeners.items():
-
             if listener.send_stop:
                 await self._stop_listener(query_id)
                 listener.send_stop = False
@@ -556,7 +562,7 @@ async def wait_closed(self) -> None:
         try:
             await asyncio.wait_for(self._wait_closed.wait(), self.close_timeout)
         except asyncio.TimeoutError:
-            log.debug("Timer close_timeout fired in wait_closed")
+            log.warning("Timer close_timeout fired in wait_closed")
 
         log.debug("wait_close: done")
 

tests/conftest.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import pathlib
+import platform
 import re
 import ssl
 import sys
@@ -19,6 +20,9 @@
 all_transport_dependencies = ["aiohttp", "requests", "httpx", "websockets", "botocore"]
 
 
+PyPy = platform.python_implementation() == "PyPy"
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--run-online",

tests/test_aiohttp_websocket_graphqlws_subscription.py

@@ -10,7 +10,7 @@
 from gql import Client, gql
 from gql.transport.exceptions import TransportConnectionClosed, TransportServerError
 
-from .conftest import MS, WebSocketServerHelper
+from .conftest import MS, PyPy, WebSocketServerHelper
 
 # Marking all tests in this file with the aiohttp AND websockets marker
 pytestmark = [pytest.mark.aiohttp, pytest.mark.websockets]
@@ -260,7 +260,8 @@ async def test_aiohttp_websocket_graphqlws_subscription_break(
     count = 10
     subscription = gql(subscription_str.format(count=count))
 
-    async for result in session.subscribe(subscription):
+    generator = session.subscribe(subscription)
+    async for result in generator:
 
         number = result["number"]
         print(f"Number received: {number}")
@@ -274,6 +275,9 @@ async def test_aiohttp_websocket_graphqlws_subscription_break(
 
     assert count == 5
 
+    # Using aclose here to make it stop cleanly on pypy
+    await generator.aclose()
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("graphqlws_server", [server_countdown], indirect=True)
@@ -847,30 +851,42 @@ async def test_aiohttp_websocket_graphqlws_subscription_reconnecting_session(
     # Then with the same session handle, we make a subscription or an execute
     # which will detect that the transport is closed so that the client could
     # try to reconnect
+    generator = None
     try:
         if execute_instead_of_subscribe:
             print("\nEXECUTION_2\n")
             await session.execute(subscription)
         else:
             print("\nSUBSCRIPTION_2\n")
-            async for result in session.subscribe(subscription):
+            generator = session.subscribe(subscription)
+            async for result in generator:
                 pass
-    except TransportClosed:
+    except (TransportClosed, TransportConnectionClosed):
+        if generator:
+            await generator.aclose()
         pass
 
-    await asyncio.sleep(50 * MS)
+    timeout = 50
+
+    if PyPy:
+        timeout = 500
+
+    await asyncio.sleep(timeout * MS)
 
     # And finally with the same session handle, we make a subscription
     # which works correctly
     print("\nSUBSCRIPTION_3\n")
-    async for result in session.subscribe(subscription):
+    generator = session.subscribe(subscription)
+    async for result in generator:
 
         number = result["number"]
         print(f"Number received: {number}")
 
         assert number == count
         count -= 1
 
+    await generator.aclose()
+
     assert count == -1
 
     await client.close_async()

tests/test_aiohttp_websocket_subscription.py

@@ -250,7 +250,8 @@ async def test_aiohttp_websocket_subscription_break(
     count = 10
     subscription = gql(subscription_str.format(count=count))
 
-    async for result in session.subscribe(subscription):
+    generator = session.subscribe(subscription)
+    async for result in generator:
 
         number = result["number"]
         print(f"Number received: {number}")
@@ -264,6 +265,9 @@ async def test_aiohttp_websocket_subscription_break(
 
     assert count == 5
 
+    # Using aclose here to make it stop cleanly on pypy
+    await generator.aclose()
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("server", [server_countdown], indirect=True)

tests/test_client.py

@@ -280,3 +280,7 @@ async def test_async_transport_close_on_schema_retrieval_failure():
         pass
 
     assert client.transport.session is None
+
+    import asyncio
+
+    await asyncio.sleep(1)

tests/test_graphqlws_subscription.py

@@ -10,7 +10,7 @@
 from gql import Client, gql
 from gql.transport.exceptions import TransportConnectionClosed, TransportServerError
 
-from .conftest import MS, WebSocketServerHelper
+from .conftest import MS, PyPy, WebSocketServerHelper
 
 # Marking all tests in this file with the websockets marker
 pytestmark = pytest.mark.websockets
@@ -260,7 +260,8 @@ async def test_graphqlws_subscription_break(
     count = 10
     subscription = gql(subscription_str.format(count=count))
 
-    async for result in session.subscribe(subscription):
+    generator = session.subscribe(subscription)
+    async for result in generator:
 
         number = result["number"]
         print(f"Number received: {number}")
@@ -274,6 +275,9 @@ async def test_graphqlws_subscription_break(
 
     assert count == 5
 
+    # Using aclose here to make it stop cleanly on pypy
+    await generator.aclose()
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("graphqlws_server", [server_countdown], indirect=True)
@@ -843,30 +847,42 @@ async def test_graphqlws_subscription_reconnecting_session(
     # Then with the same session handle, we make a subscription or an execute
     # which will detect that the transport is closed so that the client could
     # try to reconnect
+    generator = None
     try:
         if execute_instead_of_subscribe:
             print("\nEXECUTION_2\n")
             await session.execute(subscription)
         else:
             print("\nSUBSCRIPTION_2\n")
-            async for result in session.subscribe(subscription):
+            generator = session.subscribe(subscription)
+            async for result in generator:
                 pass
-    except TransportClosed:
+    except (TransportClosed, TransportConnectionClosed):
+        if generator:
+            await generator.aclose()
         pass
 
-    await asyncio.sleep(50 * MS)
+    timeout = 50
+
+    if PyPy:
+        timeout = 500
+
+    await asyncio.sleep(timeout * MS)
 
     # And finally with the same session handle, we make a subscription
     # which works correctly
     print("\nSUBSCRIPTION_3\n")
-    async for result in session.subscribe(subscription):
+    generator = session.subscribe(subscription)
+    async for result in generator:
 
         number = result["number"]
         print(f"Number received: {number}")
 
         assert number == count
         count -= 1
 
+    await generator.aclose()
+
     assert count == -1
 
     await client.close_async()

tests/test_phoenix_channel_query.py

@@ -216,8 +216,12 @@ async def test_phoenix_channel_subscription(event_loop, server, query_str):
     first_result = None
     query = gql(query_str)
     async with Client(transport=transport) as session:
-        async for result in session.subscribe(query):
+        generator = session.subscribe(query)
+        async for result in generator:
             first_result = result
             break
 
+        # Using aclose here to make it stop cleanly on pypy
+        await generator.aclose()
+
     print("Client received:", first_result)

tests/test_phoenix_channel_subscription.py

@@ -201,7 +201,9 @@ async def test_phoenix_channel_subscription(
     subscription = gql(subscription_str.format(count=count))
 
     async with Client(transport=sample_transport) as session:
-        async for result in session.subscribe(subscription):
+
+        generator = session.subscribe(subscription)
+        async for result in generator:
             number = result["countdown"]["number"]
             print(f"Number received: {number}")
 
@@ -212,6 +214,9 @@ async def test_phoenix_channel_subscription(
 
             count -= 1
 
+        # Using aclose here to make it stop cleanly on pypy
+        await generator.aclose()
+
     assert count == end_count
 
 
@@ -378,7 +383,8 @@ async def test_phoenix_channel_heartbeat(event_loop, server, subscription_str):
     subscription = gql(heartbeat_subscription_str)
     async with Client(transport=sample_transport) as session:
         i = 0
-        async for result in session.subscribe(subscription):
+        generator = session.subscribe(subscription)
+        async for result in generator:
             heartbeat_count = result["heartbeat"]["heartbeat_count"]
             print(f"Heartbeat count received: {heartbeat_count}")
 
@@ -387,3 +393,6 @@ async def test_phoenix_channel_heartbeat(event_loop, server, subscription_str):
                 break
 
             i += 1
+
+        # Using aclose here to make it stop cleanly on pypy
+        await generator.aclose()

tests/test_websocket_subscription.py

@@ -11,7 +11,7 @@
 from gql import Client, gql
 from gql.transport.exceptions import TransportConnectionClosed, TransportServerError
 
-from .conftest import MS, WebSocketServerHelper
+from .conftest import MS, PyPy, WebSocketServerHelper
 
 # Marking all tests in this file with the websockets marker
 pytestmark = pytest.mark.websockets
@@ -181,7 +181,8 @@ async def test_websocket_subscription_break(
     count = 10
     subscription = gql(subscription_str.format(count=count))
 
-    async for result in session.subscribe(subscription):
+    generator = session.subscribe(subscription)
+    async for result in generator:
 
         number = result["number"]
         print(f"Number received: {number}")
@@ -195,6 +196,9 @@ async def test_websocket_subscription_break(
 
     assert count == 5
 
+    # Using aclose here to make it stop cleanly on pypy
+    await generator.aclose()
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("server", [server_countdown], indirect=True)
@@ -413,7 +417,14 @@ async def test_websocket_subscription_with_keepalive_with_timeout_ok(
 
     path = "/graphql"
     url = f"ws://{server.hostname}:{server.port}{path}"
-    sample_transport = WebsocketsTransport(url=url, keep_alive_timeout=(20 * MS))
+
+    keep_alive_timeout = 20 * MS
+    if PyPy:
+        keep_alive_timeout = 200 * MS
+
+    sample_transport = WebsocketsTransport(
+        url=url, keep_alive_timeout=keep_alive_timeout
+    )
 
     client = Client(transport=sample_transport)