Handle connection_init and connection_terminate

graphene_subscriptions/consumers.py

@@ -49,7 +49,10 @@ def websocket_receive(self, message):
         id = request.get("id")
 
         if request["type"] == "connection_init":
-            return
+            self._send_connection_ack()
+
+        elif request["type"] == "connection_terminate":
+            self.websocket_disconnect(message)
 
         elif request["type"] == "start":
             payload = request["payload"]
@@ -95,3 +98,15 @@ def _send_result(self, id, result):
                 ),
             }
         )
+
+    def _send_connection_ack(self):
+        self.send(
+            {
+                "type": "websocket.send",
+                "text": json.dumps(
+                    {
+                        "type": "connection_ack",
+                    }
+                ),
+            }
+        )

tests/test_consumers.py

@@ -0,0 +1,31 @@
+import pytest
+from channels.testing import WebsocketCommunicator
+
+from graphene_subscriptions.consumers import GraphqlSubscriptionConsumer
+
+
+@pytest.mark.asyncio
+async def test_consumer_connection_init():
+    communicator = WebsocketCommunicator(GraphqlSubscriptionConsumer, "/graphql/")
+    connected, subprotocol = await communicator.connect()
+    assert connected
+
+    await communicator.send_json_to({"type": "connection_init"})
+
+    response = await communicator.receive_json_from()
+
+    assert response["type"] == "connection_ack"
+
+
+@pytest.mark.asyncio
+async def test_consumer_connection_terminate():
+    communicator = WebsocketCommunicator(GraphqlSubscriptionConsumer, "/graphql/")
+    connected, subprotocol = await communicator.connect()
+    assert connected
+
+    await communicator.send_json_to({"type": "connection_terminate"})
+
+    response = await communicator.receive_output()
+
+    assert response["type"] == "websocket.close"
+    assert response["code"] == 1000