Abstract ensuring async task is a future

Author: SmileyChris

graphql_ws/aiohttp.py

@@ -1,5 +1,5 @@
 import json
-from asyncio import ensure_future, shield
+from asyncio import shield
 
 from aiohttp import WSMsgType
 
@@ -45,9 +45,7 @@ async def _handle(self, ws, request_context=None):
                 break
 
             connection_context.remember_task(
-                ensure_future(
-                    self.on_message(connection_context, message), loop=self.loop
-                )
+                self.on_message(connection_context, message), loop=self.loop
             )
         await self.on_close(connection_context)
 

graphql_ws/base_async.py

@@ -35,8 +35,8 @@ def closed(self):
     async def close(self, code):
         ...
 
-    def remember_task(self, task):
-        self.pending_tasks.add(asyncio.ensure_future(task))
+    def remember_task(self, task, loop=None):
+        self.pending_tasks.add(asyncio.ensure_future(task, loop=loop))
         # Clear completed tasks
         self.pending_tasks -= WeakSet(
             task for task in self.pending_tasks if task.done()

graphql_ws/websockets_lib.py

@@ -1,5 +1,5 @@
 import json
-from asyncio import ensure_future, shield
+from asyncio import shield
 
 from websockets import ConnectionClosed
 
@@ -41,9 +41,7 @@ async def _handle(self, ws, request_context):
                 break
 
             connection_context.remember_task(
-                ensure_future(
-                    self.on_message(connection_context, message), loop=self.loop
-                )
+                self.on_message(connection_context, message), loop=self.loop
             )
         await self.on_close(connection_context)