def test_subscription_cancellation(test_client):
with test_client.websocket_connect("/", [GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
ws.send_json(ConnectionInitMessage().as_dict())
response = ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query='subscription { echo(message: "Hi", delay: 99) }'
),
).as_dict()
)
ws.send_json(
SubscribeMessage(
id="sub2",
payload=SubscribeMessagePayload(
query="subscription { debug { numActiveResultHandlers } }",
),
).as_dict()
)
response = ws.receive_json()
assert (
response
== NextMessage(
id="sub2", payload={"data": {"debug": {"numActiveResultHandlers": 2}}}
).as_dict()
)
response = ws.receive_json()
assert response == CompleteMessage(id="sub2").as_dict()
ws.send_json(CompleteMessage(id="sub1").as_dict())
ws.send_json(
SubscribeMessage(
id="sub3",
payload=SubscribeMessagePayload(
query="subscription { debug { numActiveResultHandlers } }",
),
).as_dict()
)
response = ws.receive_json()
assert (
response
== NextMessage(
id="sub3", payload={"data": {"debug": {"numActiveResultHandlers": 1}}}
).as_dict()
)
response = ws.receive_json()
assert response == CompleteMessage(id="sub3").as_dict()
ws.close()
async def test_simple_subscription(aiohttp_client):
app = create_app()
aiohttp_app_client = await aiohttp_client(app)
async with aiohttp_app_client.ws_connect(
"/graphql", protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
await ws.send_json(ConnectionInitMessage().as_dict())
response = await ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
await ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query='subscription { echo(message: "Hi") }'),
).as_dict())
response = await ws.receive_json()
assert (response == NextMessage(id="sub1",
payload={
"data": {
"echo": "Hi"
}
}).as_dict())
await ws.send_json(CompleteMessage(id="sub1").as_dict())
await ws.close()
assert ws.closed
async def test_connection_init_timeout_cancellation(test_client):
app = create_app(connection_init_wait_timeout=timedelta(milliseconds=500))
test_client = TestClient(app)
with test_client.websocket_connect("/", [GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
ws.send_json(ConnectionInitMessage().as_dict())
response = ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
await asyncio.sleep(1)
ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query="subscription { debug { isConnectionInitTimeoutTaskDone } }"
),
).as_dict()
)
response = ws.receive_json()
assert (
response
== NextMessage(
id="sub1",
payload={"data": {"debug": {"isConnectionInitTimeoutTaskDone": True}}},
).as_dict()
)
ws.close()
def test_simple_subscription(test_client):
with test_client.websocket_connect("/", [GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
ws.send_json(ConnectionInitMessage().as_dict())
response = ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query='subscription { echo(message: "Hi") }'
),
).as_dict()
)
response = ws.receive_json()
assert (
response
== NextMessage(id="sub1", payload={"data": {"echo": "Hi"}}).as_dict()
)
ws.send_json(CompleteMessage(id="sub1").as_dict())
ws.close()
def test_server_sent_ping(test_client):
with test_client.websocket_connect("/", [GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
ws.send_json(ConnectionInitMessage().as_dict())
response = ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(query="subscription { requestPing }"),
).as_dict()
)
response = ws.receive_json()
assert response == PingMessage().as_dict()
ws.send_json(PongMessage().as_dict())
response = ws.receive_json()
assert (
response
== NextMessage(id="sub1", payload={"data": {"requestPing": True}}).as_dict()
)
response = ws.receive_json()
assert response == CompleteMessage(id="sub1").as_dict()
ws.close()
async def handle_async_results(
self,
result_source: AsyncGenerator,
operation_id: str,
) -> None:
try:
async for result in result_source:
if result.errors:
error_payload = [
format_graphql_error(err) for err in result.errors
]
error_message = ErrorMessage(id=operation_id,
payload=error_payload)
await self.send_message(error_message)
self.schema.process_errors(result.errors)
return
else:
next_payload = {"data": result.data}
next_message = NextMessage(id=operation_id,
payload=next_payload)
await self.send_message(next_message)
except asyncio.CancelledError:
# CancelledErrors are expected during task cleanup.
return
except Exception as error:
# GraphQLErrors are handled by graphql-core and included in the
# ExecutionResult
error = GraphQLError(str(error), original_error=error)
error_payload = [format_graphql_error(error)]
error_message = ErrorMessage(id=operation_id,
payload=error_payload)
await self.send_message(error_message)
self.schema.process_errors([error])
return
await self.send_message(CompleteMessage(id=operation_id))
async def test_server_sent_ping(aiohttp_client):
app = create_app()
aiohttp_app_client = await aiohttp_client(app)
async with aiohttp_app_client.ws_connect(
"/graphql", protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
await ws.send_json(ConnectionInitMessage().as_dict())
response = await ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
await ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query="subscription { requestPing }"),
).as_dict())
response = await ws.receive_json()
assert response == PingMessage().as_dict()
await ws.send_json(PongMessage().as_dict())
response = await ws.receive_json()
assert (response == NextMessage(id="sub1",
payload={
"data": {
"requestPing": True
}
}).as_dict())
response = await ws.receive_json()
assert response == CompleteMessage(id="sub1").as_dict()
await ws.close()
assert ws.closed
async def test_connection_init_timeout_cancellation(aiohttp_client):
app = create_app(connection_init_wait_timeout=timedelta(milliseconds=50))
aiohttp_app_client = await aiohttp_client(app)
async with aiohttp_app_client.ws_connect(
"/graphql", protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
await ws.send_json(ConnectionInitMessage().as_dict())
response = await ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
await asyncio.sleep(0.1)
await ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query=
"subscription { debug { isConnectionInitTimeoutTaskDone } }"
),
).as_dict())
response = await ws.receive_json()
assert (response == NextMessage(
id="sub1",
payload={
"data": {
"debug": {
"isConnectionInitTimeoutTaskDone": True
}
}
},
).as_dict())
await ws.close()
assert ws.closed
async def test_subscription_cancellation(aiohttp_client):
app = create_app()
aiohttp_app_client = await aiohttp_client(app)
async with aiohttp_app_client.ws_connect(
"/graphql", protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL]) as ws:
await ws.send_json(ConnectionInitMessage().as_dict())
response = await ws.receive_json()
assert response == ConnectionAckMessage().as_dict()
await ws.send_json(
SubscribeMessage(
id="sub1",
payload=SubscribeMessagePayload(
query='subscription { echo(message: "Hi", delay: 99) }'),
).as_dict())
await ws.send_json(
SubscribeMessage(
id="sub2",
payload=SubscribeMessagePayload(
query="subscription { debug { numActiveResultHandlers } }",
),
).as_dict())
response = await ws.receive_json()
assert (response == NextMessage(id="sub2",
payload={
"data": {
"debug": {
"numActiveResultHandlers":
2
}
}
}).as_dict())
response = await ws.receive_json()
assert response == CompleteMessage(id="sub2").as_dict()
await ws.send_json(CompleteMessage(id="sub1").as_dict())
await ws.send_json(
SubscribeMessage(
id="sub3",
payload=SubscribeMessagePayload(
query="subscription { debug { numActiveResultHandlers } }",
),
).as_dict())
response = await ws.receive_json()
assert (response == NextMessage(id="sub3",
payload={
"data": {
"debug": {
"numActiveResultHandlers":
1
}
}
}).as_dict())
response = await ws.receive_json()
assert response == CompleteMessage(id="sub3").as_dict()
await ws.close()
assert ws.closed