Merge branch 'deduplicate' into channels2

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,12 +119,15 @@ 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)
-        assert message, "You need to send at least one thing"
-        return connection_context.send(message)
+        if op_id is None or connection_context.has_operation(op_id):
+            message = self.build_message(op_id, op_type, payload)
+            return connection_context.send(message)
 
     def build_message(self, id, op_type, payload):
         message = {}
@@ -131,6 +137,7 @@ def build_message(self, id, op_type, payload):
             message["type"] = op_type
         if payload is not None:
             message["payload"] = payload
+        assert message, "You need to send at least one thing"
         return message
 
     def send_execution_result(self, connection_context, op_id, execution_result):
@@ -171,11 +178,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
@@ -107,9 +121,6 @@ def process_message(self, connection_context, parsed_message):
         connection_context.remember_task(task)
         return task
 
-    async def send_message(self, *args, **kwargs):
-        await super().send_message(*args, **kwargs)
-
     async def on_open(self, connection_context):
         pass
 
@@ -125,11 +136,13 @@ 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):
-        execution_result = self.execute(params)
+        # Attempt to unsubscribe first in case we already have a subscription
+        # with this id.
+        await connection_context.unsubscribe(op_id)
 
-        if is_awaitable(execution_result):
-            execution_result = await execution_result
+        execution_result = self.execute(params)
 
+        connection_context.register_operation(op_id, execution_result)
         if hasattr(execution_result, "__aiter__"):
             iterator = await execution_result.__aiter__()
             connection_context.register_operation(op_id, iterator)
@@ -142,30 +155,25 @@ async def on_start(self, connection_context, op_id, params):
                     )
             except Exception as e:
                 await self.send_error(connection_context, op_id, e)
-            connection_context.remove_operation(op_id)
         else:
             try:
+                if is_awaitable(execution_result):
+                    execution_result = await execution_result
                 await self.send_execution_result(
                     connection_context, op_id, execution_result
                 )
             except Exception as e:
                 await self.send_error(connection_context, op_id, e)
         await self.send_message(connection_context, op_id, GQL_COMPLETE)
+        connection_context.remove_operation(op_id)
         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 send_message(
+        self, connection_context, op_id=None, op_type=None, payload=None
+    ):
+        if op_id is None or connection_context.has_operation(op_id):
+            message = self.build_message(op_id, op_type, payload)
+            return await connection_context.send(message)
 
     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):