Skip to content

Commit faadac4

Browse files
authored
websockets transport - optional variables and operation name (#111)
1 parent 140547d commit faadac4

File tree

2 files changed

+35
-9
lines changed

2 files changed

+35
-9
lines changed

gql/transport/websockets.py

+7-9
Original file line numberDiff line numberDiff line change
@@ -243,16 +243,14 @@ async def _send_query(
243243
query_id = self.next_query_id
244244
self.next_query_id += 1
245245

246+
payload: Dict[str, Any] = {"query": print_ast(document)}
247+
if variable_values:
248+
payload["variables"] = variable_values
249+
if operation_name:
250+
payload["operationName"] = operation_name
251+
246252
query_str = json.dumps(
247-
{
248-
"id": str(query_id),
249-
"type": "start",
250-
"payload": {
251-
"variables": variable_values or {},
252-
"operationName": operation_name or "",
253-
"query": print_ast(document),
254-
},
255-
}
253+
{"id": str(query_id), "type": "start", "payload": payload}
256254
)
257255

258256
await self._send(query_str)

tests/test_websocket_subscription.py

+28
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,34 @@ async def test_websocket_subscription_slow_consumer(
320320
assert count == -1
321321

322322

323+
@pytest.mark.asyncio
324+
@pytest.mark.parametrize("server", [server_countdown], indirect=True)
325+
@pytest.mark.parametrize("subscription_str", [countdown_subscription_str])
326+
async def test_websocket_subscription_with_operation_name(
327+
event_loop, client_and_server, subscription_str
328+
):
329+
330+
session, server = client_and_server
331+
332+
count = 10
333+
subscription = gql(subscription_str.format(count=count))
334+
335+
async for result in session.subscribe(
336+
subscription, operation_name="CountdownSubscription"
337+
):
338+
339+
number = result["number"]
340+
print(f"Number received: {number}")
341+
342+
assert number == count
343+
count -= 1
344+
345+
assert count == -1
346+
347+
# Check that the query contains the operationName
348+
assert '"operationName": "CountdownSubscription"' in logged_messages[0]
349+
350+
323351
WITH_KEEPALIVE = True
324352

325353

0 commit comments

Comments
 (0)