Using SubscriptionTransportBase instead of WebsocketsTransportBase for AppSync transport

Author: leszekhanusz

gql/transport/appsync_websockets.py

@@ -7,8 +7,10 @@
 from graphql import DocumentNode, ExecutionResult, print_ast
 
 from .appsync_auth import AppSyncAuthentication, AppSyncIAMAuthentication
+from .common.adapters.websockets import WebSocketsAdapter
+from .common.base import SubscriptionTransportBase
 from .exceptions import TransportProtocolError, TransportServerError
-from .websockets import WebsocketsTransport, WebsocketsTransportBase
+from .websockets import WebsocketsTransport
 
 log = logging.getLogger("gql.transport.appsync")
 
@@ -19,7 +21,7 @@
     pass
 
 
-class AppSyncWebsocketsTransport(WebsocketsTransportBase):
+class AppSyncWebsocketsTransport(SubscriptionTransportBase):
     """:ref:`Async Transport <async_transports>` used to execute GraphQL subscription on
     AWS appsync realtime endpoint.
 
@@ -32,6 +34,7 @@ class AppSyncWebsocketsTransport(WebsocketsTransportBase):
     def __init__(
         self,
         url: str,
+        *,
         auth: Optional[AppSyncAuthentication] = None,
         session: Optional["botocore.session.Session"] = None,
         ssl: Union[SSLContext, bool] = False,
@@ -70,17 +73,25 @@ def __init__(
             auth = AppSyncIAMAuthentication(host=host, session=session)
 
         self.auth = auth
+        self.ack_timeout: Optional[Union[int, float]] = ack_timeout
+        self.init_payload: Dict[str, Any] = {}
 
         url = self.auth.get_auth_url(url)
 
-        super().__init__(
-            url,
+        # Instanciate a WebSocketAdapter to indicate the use
+        # of the websockets dependency for this transport
+        self.adapter: WebSocketsAdapter = WebSocketsAdapter(
+            url=url,
             ssl=ssl,
+            connect_args=connect_args,
+        )
+
+        # Initialize the generic SubscriptionTransportBase parent class
+        super().__init__(
+            adapter=self.adapter,
             connect_timeout=connect_timeout,
             close_timeout=close_timeout,
-            ack_timeout=ack_timeout,
             keep_alive_timeout=keep_alive_timeout,
-            connect_args=connect_args,
         )
 
         # Using the same 'graphql-ws' protocol as the apollo protocol
@@ -181,7 +192,7 @@ async def _send_query(
 
         return query_id
 
-    subscribe = WebsocketsTransportBase.subscribe  # type: ignore[assignment]
+    subscribe = SubscriptionTransportBase.subscribe  # type: ignore[assignment]
     """Send a subscription query and receive the results using
     a python async generator.
 
@@ -212,3 +223,19 @@ async def execute(
         WebsocketsTransport._send_init_message_and_wait_ack
     )
     _wait_ack = WebsocketsTransport._wait_ack
+
+    @property
+    def url(self) -> str:
+        return self.adapter.url
+
+    @property
+    def headers(self) -> Dict[str, str]:
+        return self.adapter.headers
+
+    @property
+    def ssl(self) -> Union[SSLContext, bool]:
+        return self.adapter.ssl
+
+    @property
+    def connect_args(self) -> Dict[str, Any]:
+        return self.adapter.connect_args