Merge pull request #36 from Pelleplutt/master

Add ability to override init_payload for subscriptions.

Author: DaleSeo

python_graphql_client/graphql_client.py

@@ -82,9 +82,12 @@ async def subscribe(
         variables: dict = None,
         operation_name: str = None,
         headers: dict = {},
+        init_payload: dict = {},
     ):
         """Make asynchronous request for GraphQL subscription."""
-        connection_init_message = json.dumps({"type": "connection_init", "payload": {}})
+        connection_init_message = json.dumps(
+            {"type": "connection_init", "payload": init_payload}
+        )
 
         request_body = self.__request_body(
             query=query, variables=variables, operation_name=operation_name

tests/test_graphql_client.py

@@ -323,3 +323,40 @@ async def test_headers_passed_to_websocket_connect(self, mock_connect):
         )
 
         mock_handle.assert_has_calls([call({"data": {"messageAdded": "one"}})])
+
+    @patch("websockets.connect")
+    async def test_init_payload_passed_in_init_message(self, mock_connect):
+        """Subsribe a GraphQL subscription."""
+        mock_websocket = mock_connect.return_value.__aenter__.return_value
+        mock_websocket.send = AsyncMock()
+        mock_websocket.__aiter__.return_value = [
+            '{"type": "connection_init", "payload": '
+            '{"init": "this is the init_payload"}}',
+            '{"type": "data", "id": "1", "payload": {"data": {"messageAdded": "one"}}}',
+        ]
+        expected_endpoint = "ws://www.test-api.com/graphql"
+        client = GraphqlClient(endpoint=expected_endpoint)
+
+        query = """
+            subscription onMessageAdded {
+                messageAdded
+            }
+        """
+        init_payload = '{"init": "this is the init_payload"}'
+
+        mock_handle = MagicMock()
+
+        await client.subscribe(
+            query=query, handle=mock_handle, init_payload=init_payload
+        )
+
+        mock_connect.assert_called_with(
+            expected_endpoint, subprotocols=["graphql-ws"], extra_headers={}
+        )
+
+        mock_handle.assert_has_calls(
+            [
+                call({"init": "this is the init_payload"}),
+                call({"data": {"messageAdded": "one"}}),
+            ]
+        )