async def test_subscription(
client, headers, subscription_query, mutation_city, city_name
):
request = GraphQLRequest(query=subscription_query, headers=headers)
m = []
def callback(data):
assert "city" in data
m.append(data)
if len(m) > 1:
city = data.get("city")[0]
assert city.get("name") == city_name
subscription.unsubscribe()
callbacks = CallbackRegistry()
callbacks.register(
GraphQLSubscriptionEventType.DATA, lambda event: callback(event.payload.data)
)
subscription: GraphQLSubscription = await client.subscribe(
request=request, callbacks=callbacks, headers=headers,
)
await asyncio.sleep(0.1)
request = GraphQLRequest(query=mutation_city, headers=headers)
_ = await client.query(request)
try:
await asyncio.wait_for(subscription.task, timeout=1)
assert len(m) == 2
except asyncio.TimeoutError:
pytest.fail("Subscriptions timed out before receiving expected messages")
def test_callback_registry_coro_without_loop(caplog):
async def hello():
pass
registry = CallbackRegistry(callbacks={"hello": hello})
with pytest.warns(RuntimeWarning):
registry.dispatch("hello")
assert "Callback triggered without a running loop" in caplog.text
assert f"skipping {str(hello)}" in caplog.text
def test_callback_registry_multiple(mocker):
callback_one = mocker.Mock()
callback_two = mocker.Mock()
registry = CallbackRegistry(
callbacks={"event": [callback_one, callback_two]})
assert len(registry.callbacks("event")) == 2
registry.dispatch("event", "event")
callback_one.assert_called_once()
callback_one.assert_called_once_with("event")
callback_two.assert_called_once()
callback_two.assert_called_once_with("event")
callback_one.reset_mock()
callback_two.reset_mock()
registry.deregister("event", callback_two)
registry.dispatch("event", "event")
callback_one.assert_called_once()
callback_one.assert_called_once_with("event")
callback_two.assert_not_called()
def __post_init__(
self,
headers: Optional[Dict[str, str]] = None,
operation: Optional[str] = None,
variables: Optional[Dict[str, Any]] = None,
):
super().__post_init__(headers, operation, variables)
if self.callbacks is None:
object.__setattr__(self, "callbacks", CallbackRegistry())
elif isinstance(self.callbacks, dict):
object.__setattr__(self, "callbacks",
CallbackRegistry(callbacks=self.callbacks))
def test_callback_registry_exists(mocker):
callback = mocker.Mock()
registry = CallbackRegistry()
registry.register("one", callback)
assert registry.exists("one", callback)
assert registry.exists("one", Callback(callback))
assert not registry.exists(None, callback)
assert not registry.exists("two", callback)
def test_callback_registry_callbacks(mocker):
callback_one = mocker.Mock()
callback_two = mocker.Mock()
callback_default = mocker.Mock()
registry = CallbackRegistry()
registry.register("one", callback_one)
registry.register("two", callback_two)
registry.register(None, callback_default)
assert registry.callbacks() == [Callback(callback_default)]
assert registry.callbacks("one") == [Callback(callback_one)]
assert registry.callbacks("two") == [Callback(callback_two)]
def test_callback_registry_simple(mocker):
callback_one = mocker.Mock()
callback_two = mocker.Mock()
registry = CallbackRegistry(callbacks={
"one": callback_one,
"two": callback_two
})
registry.dispatch("one", "one")
registry.dispatch("two", "two")
registry.dispatch(None, "None")
callback_one.assert_called_once()
callback_one.assert_called_once_with("one")
callback_two.assert_called_once()
callback_two.assert_called_once_with("two")
callback_one.reset_mock()
registry.deregister("one", callback_one)
registry.dispatch("one", "one")
callback_one.assert_not_called()
def test_callback_registry_callbacks_default(mocker):
callback = mocker.Mock()
registry = CallbackRegistry()
assert registry.callbacks() == []
registry.register(None, callback)
assert registry.callbacks() == [Callback(callback)]
async def subscribe(
self,
request: GraphQLRequest,
callbacks: Optional[CallbackRegistry] = None,
headers: Optional[Dict[str, str]] = None,
) -> GraphQLSubscription:
await self._validate(request, headers=headers)
headers = headers or {}
subscription = GraphQLSubscription(
request=request,
callbacks=callbacks or CallbackRegistry(),
headers={**self._headers, **request.headers, **headers},
)
subscription.task = asyncio.create_task(self._subscribe(subscription))
return subscription
async def subscribe(
self,
request: GraphQLRequest,
headers: Optional[Dict[str, str]] = None,
operation: Optional[str] = None,
variables: Optional[Dict[str, Any]] = None,
callbacks: Optional[CallbacksType] = None,
on_data: Optional[CallbackType] = None,
on_error: Optional[CallbackType] = None,
session: Optional[aiohttp.ClientSession] = None,
wait: bool = False,
) -> GraphQLSubscription:
"""
Create and initialise a GraphQL subscription. Once subscribed and a known event
is received, all registered callbacks for the event type is triggered with the
:class:`aiographql.client.GraphQLSubscriptionEvent` instance passed in the first
argument.
The following example will start a subscription that prints all data events as
it receives them.
.. code-block:: python
# initialise and subscribe to events in the background
subscription: GraphQLSubscription = await client.subscribe(
request="{ notifications: { id, summary } }",
on_data=lambda event: print(f"Data: {event}"),
on_error=lambda event: print(f"Error: {event}"),
)
# process events for 10 seconds then unsubscribe
await asyncio.wait(subscription.task, timeout=10)
subscription.unsubscribe()
:param request: Request to send to the GraphQL server.
:param headers: Additional headers to be set when sending HTTP request.
:param operation: GraphQL operation name to use if the `GraphQLRequest.query`
contains named operations. This will override any default
operation set.
:param variables: Query variables to set for the provided request. This will
override the default values for any existing variables in the
request if set.
:param session: Optional `aiohttp.ClientSession` to use for requests
:return: The resulting `GraphQLResponse` object.
:param callbacks: Custom callback registry mapping an event to one more more
callback methods. If not provided, a new instance is created.
:param on_data: Callback to use when data event is received.
:param on_error: Callback to use when an error occurs.
:param session: Optional session to use for connecting the graphql endpoint, if
one is not provided, a new session is created for the duration of the
subscription.
:param wait: If set to `True`, this method will wait until the subscription
is completed, websocket disconnected or async task cancelled.
:return: The initialised subscription.
"""
request = self._prepare_request(request=request,
operation=operation,
variables=variables,
headers=headers)
await self.validate(request=request)
callbacks = callbacks or CallbackRegistry()
if on_data:
callbacks.register(GraphQLSubscriptionEventType.DATA, on_data)
if on_error:
callbacks.register(GraphQLSubscriptionEventType.ERROR, on_error)
subscription = GraphQLSubscription(request=request,
callbacks=callbacks)
await subscription.subscribe(endpoint=self.endpoint,
session=session or self._session,
wait=wait)
return subscription
def test_callback_registry_deregister_non_existing():
# noinspection PyBroadException
try:
CallbackRegistry().deregister("one", lambda x: None)
except Exception:
pytest.fail("Unexpected exception")
def test_callback_registry_default(mocker):
callback = mocker.Mock()
registry = CallbackRegistry()
registry.register(None, callback)
registry.dispatch("doesnotexist", "one")
callback.assert_called_once()
callback.assert_called_once_with("one")
callback.reset_mock()
registry.dispatch(None, "one")
callback.assert_called_once()
callback.assert_called_once_with("one")
callback.reset_mock()
registry.deregister(None, callback)
registry.dispatch("doesnotexist", "two")
registry.dispatch(None, "two")
callback.assert_not_called()
