Merge pull request #536 from leszekhanusz/chore_refactor_websockets

Refactor websockets transports

Author: leszekhanusz

gql/transport/aiohttp.py

@@ -1,5 +1,4 @@
 import asyncio
-import functools
 import io
 import json
 import logging
@@ -28,6 +27,7 @@
 from ..utils import extract_files
 from .appsync_auth import AppSyncAuthentication
 from .async_transport import AsyncTransport
+from .common.aiohttp_closed_event import create_aiohttp_closed_event
 from .exceptions import (
     TransportAlreadyConnected,
     TransportClosed,
@@ -147,59 +147,6 @@ async def connect(self) -> None:
         else:
             raise TransportAlreadyConnected("Transport is already connected")
 
-    @staticmethod
-    def create_aiohttp_closed_event(session) -> asyncio.Event:
-        """Work around aiohttp issue that doesn't properly close transports on exit.
-
-        See https://github.com/aio-libs/aiohttp/issues/1925#issuecomment-639080209
-
-        Returns:
-           An event that will be set once all transports have been properly closed.
-        """
-
-        ssl_transports = 0
-        all_is_lost = asyncio.Event()
-
-        def connection_lost(exc, orig_lost):
-            nonlocal ssl_transports
-
-            try:
-                orig_lost(exc)
-            finally:
-                ssl_transports -= 1
-                if ssl_transports == 0:
-                    all_is_lost.set()
-
-        def eof_received(orig_eof_received):
-            try:  # pragma: no cover
-                orig_eof_received()
-            except AttributeError:  # pragma: no cover
-                # It may happen that eof_received() is called after
-                # _app_protocol and _transport are set to None.
-                pass
-
-        for conn in session.connector._conns.values():
-            for handler, _ in conn:
-                proto = getattr(handler.transport, "_ssl_protocol", None)
-                if proto is None:
-                    continue
-
-                ssl_transports += 1
-                orig_lost = proto.connection_lost
-                orig_eof_received = proto.eof_received
-
-                proto.connection_lost = functools.partial(
-                    connection_lost, orig_lost=orig_lost
-                )
-                proto.eof_received = functools.partial(
-                    eof_received, orig_eof_received=orig_eof_received
-                )
-
-        if ssl_transports == 0:
-            all_is_lost.set()
-
-        return all_is_lost
-
     async def close(self) -> None:
         """Coroutine which will close the aiohttp session.
 
@@ -219,7 +166,7 @@ async def close(self) -> None:
                 log.debug("connector_owner is False -> not closing connector")
 
             else:
-                closed_event = self.create_aiohttp_closed_event(self.session)
+                closed_event = create_aiohttp_closed_event(self.session)
                 await self.session.close()
                 try:
                     await asyncio.wait_for(closed_event.wait(), self.ssl_close_timeout)