Move unsubscribe logic to the connection context

Author: SmileyChris

graphql_ws/base.py

@@ -35,9 +35,18 @@ def get_operation(self, op_id):
 
     def remove_operation(self, op_id):
         try:
-            del self.operations[op_id]
+            return self.operations.pop(op_id)
         except KeyError:
-            pass
+            return
+
+    def unsubscribe(self, op_id):
+        async_iterator = self.remove_operation(op_id)
+        if hasattr(async_iterator, 'dispose'):
+            async_iterator.dispose()
+
+    def unsubscribe_all(self):
+        for op_id in list(self.operations):
+            self.unsubscribe(op_id)
 
     def receive(self):
         raise NotImplementedError("receive method not implemented")
@@ -76,12 +85,6 @@ def process_message(self, connection_context, parsed_message):
 
         elif op_type == GQL_START:
             assert isinstance(payload, dict), "The payload must be a dict"
-
-            # If we already have a subscription with this id, unsubscribe from
-            # it first
-            if connection_context.has_operation(op_id):
-                self.unsubscribe(connection_context, op_id)
-
             params = self.get_graphql_params(connection_context, payload)
             return self.on_start(connection_context, op_id, params)
 
@@ -116,7 +119,10 @@ def on_open(self, connection_context):
         raise NotImplementedError("on_open method not implemented")
 
     def on_stop(self, connection_context, op_id):
-        raise NotImplementedError("on_stop method not implemented")
+        return connection_context.unsubscribe(op_id)
+
+    def on_close(self, connection_context):
+        return connection_context.unsubscribe_all()
 
     def send_message(self, connection_context, op_id=None, op_type=None, payload=None):
         message = self.build_message(op_id, op_type, payload)
@@ -171,11 +177,3 @@ def on_message(self, connection_context, message):
             return self.send_error(connection_context, None, e)
 
         return self.process_message(connection_context, parsed_message)
-
-    def unsubscribe(self, connection_context, op_id):
-        if connection_context.has_operation(op_id):
-            # Close async iterator
-            connection_context.get_operation(op_id).dispose()
-            # Close operation
-            connection_context.remove_operation(op_id)
-        return self.on_operation_complete(connection_context, op_id)

graphql_ws/base_async.py

@@ -88,6 +88,20 @@ def remember_task(self, task):
             task for task in self.pending_tasks if task.done()
         )
 
+    async def unsubscribe(self, op_id):
+        super().unsubscribe(op_id)
+
+    async def unsubscribe_all(self):
+        awaitables = [self.unsubscribe(op_id) for op_id in list(self.operations)]
+        for task in self.pending_tasks:
+            task.cancel()
+            awaitables.append(task)
+        if awaitables:
+            try:
+                await asyncio.gather(*awaitables)
+            except asyncio.CancelledError:
+                pass
+
 
 class BaseAsyncSubscriptionServer(base.BaseSubscriptionServer, ABC):
     graphql_executor = AsyncioExecutor
@@ -125,6 +139,10 @@ async def on_connection_init(self, connection_context, op_id, payload):
             await connection_context.close(1011)
 
     async def on_start(self, connection_context, op_id, params):
+        # Attempt to unsubscribe first in case we already have a subscription
+        # with this id.
+        await connection_context.unsubscribe(op_id)
+
         execution_result = self.execute(params)
 
         if is_awaitable(execution_result):
@@ -153,20 +171,6 @@ async def on_start(self, connection_context, op_id, params):
         await self.send_message(connection_context, op_id, GQL_COMPLETE)
         await self.on_operation_complete(connection_context, op_id)
 
-    async def on_close(self, connection_context):
-        awaitables = tuple(
-            self.unsubscribe(connection_context, op_id)
-            for op_id in connection_context.operations
-        ) + tuple(task.cancel() for task in connection_context.pending_tasks)
-        if awaitables:
-            try:
-                await asyncio.gather(*awaitables, loop=self.loop)
-            except asyncio.CancelledError:
-                pass
-
-    async def on_stop(self, connection_context, op_id):
-        await self.unsubscribe(connection_context, op_id)
-
     async def on_operation_complete(self, connection_context, op_id):
         pass
 

graphql_ws/base_sync.py

@@ -20,11 +20,6 @@ def on_open(self, connection_context):
     def on_connect(self, connection_context, payload):
         pass
 
-    def on_close(self, connection_context):
-        remove_operations = list(connection_context.operations)
-        for op_id in remove_operations:
-            self.unsubscribe(connection_context, op_id)
-
     def on_connection_init(self, connection_context, op_id, payload):
         try:
             self.on_connect(connection_context, payload)
@@ -34,10 +29,10 @@ def on_connection_init(self, connection_context, op_id, payload):
             self.send_error(connection_context, op_id, e, GQL_CONNECTION_ERROR)
             connection_context.close(1011)
 
-    def on_stop(self, connection_context, op_id):
-        self.unsubscribe(connection_context, op_id)
-
     def on_start(self, connection_context, op_id, params):
+        # Attempt to unsubscribe first in case we already have a subscription
+        # with this id.
+        connection_context.unsubscribe(op_id)
         try:
             execution_result = self.execute(params)
             assert isinstance(

tests/test_base.py

@@ -16,8 +16,13 @@ def test_not_implemented():
         server.on_connection_init(connection_context=None, op_id=1, payload={})
     with pytest.raises(NotImplementedError):
         server.on_open(connection_context=None)
-    with pytest.raises(NotImplementedError):
-        server.on_stop(connection_context=None, op_id=1)
+
+
+def test_on_stop():
+    server = base.BaseSubscriptionServer(schema=None)
+    context = mock.Mock()
+    server.on_stop(connection_context=context, op_id=1)
+    context.unsubscribe.assert_called_with(1)
 
 
 def test_terminate():

tests/test_graphql_ws.py

@@ -94,12 +94,12 @@ def test_start_existing_op(self, ss, cc):
         ss.get_graphql_params.return_value = {"params": True}
         cc.has_operation = mock.Mock()
         cc.has_operation.return_value = True
-        ss.unsubscribe = mock.Mock()
+        cc.unsubscribe = mock.Mock()
         ss.on_start = mock.Mock()
         ss.process_message(
             cc, {"id": "1", "type": constants.GQL_START, "payload": {"a": "b"}}
         )
-        assert ss.unsubscribe.called
+        assert cc.unsubscribe.called
         ss.on_start.assert_called_with(cc, "1", {"params": True})
 
     def test_start_bad_graphql_params(self, ss, cc):
@@ -162,7 +162,8 @@ def test_build_message_partial(ss):
     assert ss.build_message(id=None, op_type=None, payload="PAYLOAD") == {
         "payload": "PAYLOAD"
     }
-    assert ss.build_message(id=None, op_type=None, payload=None) == {}
+    with pytest.raises(AssertionError):
+        ss.build_message(id=None, op_type=None, payload=None)
 
 
 def test_send_execution_result(ss):