import asyncio from typing import Any, Callable, Dict, List from graphql.execution import ( ExecutionResult, MapAsyncIterator, create_source_event_stream, subscribe, ) from graphql.language import parse from graphql.pyutils import SimplePubSub from graphql.type import ( GraphQLArgument, GraphQLBoolean, GraphQLField, GraphQLInt, GraphQLList, GraphQLObjectType, GraphQLSchema, GraphQLString, ) from pytest import mark, raises try: anext # type: ignore except NameError: # pragma: no cover (Python < 3.10) # noinspection PyShadowingBuiltins async def anext(iterator): """Return the next item from an async iterator.""" return await iterator.__anext__() Email = Dict # should become a TypedDict once we require Python 3.8 EmailType = GraphQLObjectType( "Email", { "from": GraphQLField(GraphQLString), "subject": GraphQLField(GraphQLString), "message": GraphQLField(GraphQLString), "unread": GraphQLField(GraphQLBoolean), }, ) InboxType = GraphQLObjectType( "Inbox", { "total": GraphQLField( GraphQLInt, resolve=lambda inbox, _info: len(inbox["emails"]) ), "unread": GraphQLField( GraphQLInt, resolve=lambda inbox, _info: sum( 1 for email in inbox["emails"] if email["unread"] ), ), "emails": GraphQLField(GraphQLList(EmailType)), }, ) QueryType = GraphQLObjectType("Query", {"inbox": GraphQLField(InboxType)}) EmailEventType = GraphQLObjectType( "EmailEvent", {"email": GraphQLField(EmailType), "inbox": GraphQLField(InboxType)} ) email_schema = GraphQLSchema( query=QueryType, subscription=GraphQLObjectType( "Subscription", { "importantEmail": GraphQLField( EmailEventType, args={"priority": GraphQLArgument(GraphQLInt)}, ) }, ), ) def create_subscription(pubsub: SimplePubSub): document = parse( """ subscription ($priority: Int = 0) { importantEmail(priority: $priority) { email { from subject } inbox { unread total } } } """ ) emails: List[Email] = [ { "from": "joe@graphql.org", "subject": "Hello", "message": "Hello World", "unread": False, } ] def transform(new_email): emails.append(new_email) return {"importantEmail": {"email": new_email, "inbox": data["inbox"]}} data: Dict[str, Any] = { "inbox": {"emails": emails}, "importantEmail": pubsub.get_subscriber(transform), } return subscribe(email_schema, document, data) DummyQueryType = GraphQLObjectType("Query", {"dummy": GraphQLField(GraphQLString)}) # Check all error cases when initializing the subscription. def describe_subscription_initialization_phase(): @mark.asyncio async def accepts_positional_arguments(): document = parse( """ subscription { importantEmail } """ ) async def empty_async_iterator(_info): for value in (): # type: ignore yield value # pragma: no cover ai = await subscribe( email_schema, document, {"importantEmail": empty_async_iterator} ) with raises(StopAsyncIteration): await anext(ai) await ai.aclose() # type: ignore @mark.asyncio async def accepts_multiple_subscription_fields_defined_in_schema(): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", { "foo": GraphQLField(GraphQLString), "bar": GraphQLField(GraphQLString), }, ), ) async def foo_generator(_info): yield {"foo": "FooValue"} subscription = await subscribe( schema, parse("subscription { foo }"), {"foo": foo_generator} ) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ({"foo": "FooValue"}, None) await subscription.aclose() @mark.asyncio async def accepts_type_definition_with_sync_subscribe_function(): async def foo_generator(_obj, _info): yield {"foo": "FooValue"} schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString, subscribe=foo_generator)}, ), ) subscription = await subscribe(schema, parse("subscription { foo }")) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ({"foo": "FooValue"}, None) await subscription.aclose() @mark.asyncio async def accepts_type_definition_with_async_subscribe_function(): async def foo_generator(_obj, _info): await asyncio.sleep(0) yield {"foo": "FooValue"} schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString, subscribe=foo_generator)}, ), ) subscription = await subscribe(schema, parse("subscription { foo }")) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ({"foo": "FooValue"}, None) await subscription.aclose() @mark.asyncio async def uses_a_custom_default_subscribe_field_resolver(): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString)} ), ) class Root: @staticmethod async def custom_foo(): yield {"foo": "FooValue"} subscription = await subscribe( schema, document=parse("subscription { foo }"), root_value=Root(), subscribe_field_resolver=lambda root, _info: root.custom_foo(), ) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ( {"foo": "FooValue"}, None, ) await subscription.aclose() @mark.asyncio async def should_only_resolve_the_first_field_of_invalid_multi_field(): did_resolve = {"foo": False, "bar": False} async def subscribe_foo(_obj, _info): did_resolve["foo"] = True yield {"foo": "FooValue"} async def subscribe_bar(_obj, _info): # pragma: no cover did_resolve["bar"] = True yield {"bar": "BarValue"} schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", { "foo": GraphQLField(GraphQLString, subscribe=subscribe_foo), "bar": GraphQLField(GraphQLString, subscribe=subscribe_bar), }, ), ) subscription = await subscribe(schema, parse("subscription { foo bar }")) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ( {"foo": "FooValue", "bar": None}, None, ) assert did_resolve == {"foo": True, "bar": False} await subscription.aclose() @mark.asyncio async def throws_an_error_if_some_of_required_arguments_are_missing(): document = parse("subscription { foo }") schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString)} ), ) with raises(TypeError, match="^Expected None to be a GraphQL schema\\.$"): await subscribe(None, document) # type: ignore with raises(TypeError, match="missing .* positional argument: 'schema'"): await subscribe(document=document) # type: ignore with raises(TypeError, match="^Must provide document\\.$"): await subscribe(schema, None) # type: ignore with raises(TypeError, match="missing .* positional argument: 'document'"): await subscribe(schema=schema) # type: ignore @mark.asyncio async def resolves_to_an_error_if_schema_does_not_support_subscriptions(): schema = GraphQLSchema(query=DummyQueryType) document = parse("subscription { unknownField }") result = await subscribe(schema, document) assert result == ( None, [ { "message": "Schema is not configured to execute" " subscription operation.", "locations": [(1, 1)], } ], ) @mark.asyncio async def resolves_to_an_error_for_unknown_subscription_field(): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString)} ), ) document = parse("subscription { unknownField }") result = await subscribe(schema, document) assert result == ( None, [ { "message": "The subscription field 'unknownField' is not defined.", "locations": [(1, 16)], } ], ) @mark.asyncio async def should_pass_through_unexpected_errors_thrown_in_subscribe(): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString)} ), ) with raises(TypeError, match="^Must provide document\\.$"): await subscribe(schema=schema, document={}) # type: ignore @mark.asyncio @mark.filterwarnings("ignore:.* was never awaited:RuntimeWarning") async def throws_an_error_if_subscribe_does_not_return_an_iterator(): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", { "foo": GraphQLField( GraphQLString, subscribe=lambda _obj, _info: "test" ) }, ), ) document = parse("subscription { foo }") with raises(TypeError) as exc_info: await subscribe(schema, document) assert str(exc_info.value) == ( "Subscription field must return AsyncIterable. Received: 'test'." ) @mark.asyncio async def resolves_to_an_error_for_subscription_resolver_errors(): async def subscribe_with_fn(subscribe_fn: Callable): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", {"foo": GraphQLField(GraphQLString, subscribe=subscribe_fn)}, ), ) document = parse("subscription { foo }") result = await subscribe(schema, document) assert await create_source_event_stream(schema, document) == result return result expected_result = ( None, [ { "message": "test error", "locations": [(1, 16)], "path": ["foo"], } ], ) # Returning an error def return_error(_obj, _info): return TypeError("test error") assert await subscribe_with_fn(return_error) == expected_result # Throwing an error def throw_error(*_args): raise TypeError("test error") assert await subscribe_with_fn(throw_error) == expected_result # Resolving to an error async def resolve_error(*_args): return TypeError("test error") assert await subscribe_with_fn(resolve_error) == expected_result # Rejecting with an error async def reject_error(*_args): return TypeError("test error") assert await subscribe_with_fn(reject_error) == expected_result @mark.asyncio async def resolves_to_an_error_if_variables_were_wrong_type(): schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", { "foo": GraphQLField( GraphQLString, {"arg": GraphQLArgument(GraphQLInt)} ) }, ), ) variable_values = {"arg": "meow"} document = parse( """ subscription ($arg: Int) { foo(arg: $arg) } """ ) # If we receive variables that cannot be coerced correctly, subscribe() will # resolve to an ExecutionResult that contains an informative error description. result = await subscribe(schema, document, variable_values=variable_values) assert isinstance(result, ExecutionResult) assert result == ( None, [ { "message": "Variable '$arg' got invalid value 'meow';" " Int cannot represent non-integer value: 'meow'", "locations": [(2, 27)], } ], ) errors = result.errors assert errors assert errors[0].original_error # Once a subscription returns a valid AsyncIterator, it can still yield errors. def describe_subscription_publish_phase(): @mark.asyncio async def produces_a_payload_for_multiple_subscribe_in_same_subscription(): pubsub = SimplePubSub() subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) second_subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) payload1 = anext(subscription) payload2 = anext(second_subscription) assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright", "message": "Tests are good", "unread": True, } ) is True ) expected_payload = { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Alright"}, "inbox": {"unread": 1, "total": 2}, } } assert await payload1 == (expected_payload, None) assert await payload2 == (expected_payload, None) @mark.asyncio async def produces_a_payload_per_subscription_event(): pubsub = SimplePubSub() subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) # Wait for the next subscription payload. payload = anext(subscription) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright", "message": "Tests are good", "unread": True, } ) is True ) # The previously waited on payload now has a value. assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Alright"}, "inbox": {"unread": 1, "total": 2}, } }, None, ) # Another new email arrives, before anext(subscription) is called. assert ( pubsub.emit( { "from": "hyo@graphql.org", "subject": "Tools", "message": "I <3 making things", "unread": True, } ) is True ) # The next waited on payload will have a value. assert await anext(subscription) == ( { "importantEmail": { "email": {"from": "hyo@graphql.org", "subject": "Tools"}, "inbox": {"unread": 2, "total": 3}, } }, None, ) # The client decides to disconnect. # noinspection PyUnresolvedReferences await subscription.aclose() # Which may result in disconnecting upstream services as well. assert ( pubsub.emit( { "from": "adam@graphql.org", "subject": "Important", "message": "Read me please", "unread": True, } ) is False ) # No more listeners. # Awaiting subscription after closing it results in completed results. with raises(StopAsyncIteration): assert await anext(subscription) @mark.asyncio async def produces_a_payload_when_there_are_multiple_events(): pubsub = SimplePubSub() subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) payload = anext(subscription) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright", "message": "Tests are good", "unread": True, } ) is True ) assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Alright"}, "inbox": {"unread": 1, "total": 2}, } }, None, ) payload = anext(subscription) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright 2", "message": "Tests are good 2", "unread": True, } ) is True ) assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Alright 2"}, "inbox": {"unread": 2, "total": 3}, } }, None, ) @mark.asyncio async def should_not_trigger_when_subscription_is_already_done(): pubsub = SimplePubSub() subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) payload = anext(subscription) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright", "message": "Tests are good", "unread": True, } ) is True ) assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Alright"}, "inbox": {"unread": 1, "total": 2}, } }, None, ) payload = anext(subscription) await subscription.aclose() # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright 2", "message": "Tests are good 2", "unread": True, } ) is False ) with raises(StopAsyncIteration): await payload @mark.asyncio async def should_not_trigger_when_subscription_is_thrown(): pubsub = SimplePubSub() subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) payload = anext(subscription) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Alright", "message": "Tests are good", "unread": True, } ) is True ) assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Alright"}, "inbox": {"unread": 1, "total": 2}, } }, None, ) payload = anext(subscription) # Throw error with raises(RuntimeError) as exc_info: await subscription.athrow(RuntimeError("ouch")) assert str(exc_info.value) == "ouch" with raises(StopAsyncIteration): await payload @mark.asyncio async def event_order_is_correct_for_multiple_publishes(): pubsub = SimplePubSub() subscription = await create_subscription(pubsub) assert isinstance(subscription, MapAsyncIterator) payload = anext(subscription) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Message", "message": "Tests are good", "unread": True, } ) is True ) # A new email arrives! assert ( pubsub.emit( { "from": "yuzhi@graphql.org", "subject": "Message 2", "message": "Tests are good 2", "unread": True, } ) is True ) assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Message"}, "inbox": {"unread": 2, "total": 3}, } }, None, ) payload = anext(subscription) assert await payload == ( { "importantEmail": { "email": {"from": "yuzhi@graphql.org", "subject": "Message 2"}, "inbox": {"unread": 2, "total": 3}, } }, None, ) @mark.asyncio async def should_handle_error_during_execution_of_source_event(): async def generate_messages(_obj, _info): yield "Hello" yield "Goodbye" yield "Bonjour" def resolve_message(message, _info): if message == "Goodbye": raise RuntimeError("Never leave.") return message schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", { "newMessage": GraphQLField( GraphQLString, subscribe=generate_messages, resolve=resolve_message, ) }, ), ) document = parse("subscription { newMessage }") subscription = await subscribe(schema, document) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ({"newMessage": "Hello"}, None) # An error in execution is presented as such. assert await anext(subscription) == ( {"newMessage": None}, [ { "message": "Never leave.", "locations": [(1, 16)], "path": ["newMessage"], } ], ) # However that does not close the response event stream. # Subsequent events are still executed. assert await anext(subscription) == ({"newMessage": "Bonjour"}, None) @mark.asyncio async def should_pass_through_error_thrown_in_source_event_stream(): async def generate_messages(_obj, _info): yield "Hello" raise RuntimeError("test error") def resolve_message(message, _info): return message schema = GraphQLSchema( query=DummyQueryType, subscription=GraphQLObjectType( "Subscription", { "newMessage": GraphQLField( GraphQLString, resolve=resolve_message, subscribe=generate_messages, ) }, ), ) document = parse("subscription { newMessage }") subscription = await subscribe(schema, document) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ({"newMessage": "Hello"}, None) with raises(RuntimeError) as exc_info: await anext(subscription) assert str(exc_info.value) == "test error" with raises(StopAsyncIteration): await anext(subscription) @mark.asyncio async def should_work_with_async_resolve_function(): async def generate_messages(_obj, _info): yield "Hello" def resolve_message(message, _info): return message schema = GraphQLSchema( query=QueryType, subscription=GraphQLObjectType( "Subscription", { "newMessage": GraphQLField( GraphQLString, resolve=resolve_message, subscribe=generate_messages, ) }, ), ) document = parse("subscription { newMessage }") subscription = await subscribe(schema, document) assert isinstance(subscription, MapAsyncIterator) assert await anext(subscription) == ({"newMessage": "Hello"}, None)