test_subscribe.py

import asyncio
from collections.abc import AsyncIterable, AsyncIterator, Callable
from typing import (
    Any,
    TypedDict,
    TypeVar,
)

import pytest

from graphql.error import GraphQLError
from graphql.execution import (
    ExecutionResult,
    Executor,
    create_source_event_stream,
    execute_subscription_event,
    map_source_to_response_event,
    subscribe,
)
from graphql.language import DocumentNode, parse
from graphql.pyutils import AwaitableOrValue, SimplePubSub, is_awaitable
from graphql.type import (
    GraphQLArgument,
    GraphQLBoolean,
    GraphQLField,
    GraphQLInt,
    GraphQLList,
    GraphQLObjectType,
    GraphQLResolveInfo,
    GraphQLSchema,
    GraphQLString,
)

from ..fixtures import cleanup
from ..utils.assert_equal_awaitables_or_values import assert_equal_awaitables_or_values

pytestmark = [
    pytest.mark.anyio,
    pytest.mark.filterwarnings("ignore:coroutine .* was never awaited:RuntimeWarning"),
]

T = TypeVar("T")

Email = TypedDict(
    "Email",
    {
        "from": str,
        "subject": str,
        "message": str,
        "unread": bool,
    },
)


async def async_subject(email: Email, _info: GraphQLResolveInfo) -> str:
    return email["subject"]


EmailType = GraphQLObjectType(
    "Email",
    {
        "from": GraphQLField(GraphQLString),
        "subject": GraphQLField(GraphQLString),
        "asyncSubject": GraphQLField(GraphQLString, resolve=async_subject),
        "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, variable_values: dict[str, Any] | None = None
) -> AwaitableOrValue[AsyncIterator[ExecutionResult] | ExecutionResult]:
    document = parse(
        """
        subscription (
          $priority: Int = 0
          $shouldDefer: Boolean = false
          $shouldStream: Boolean = false
          $asyncResolver: Boolean = false
        ) {
          importantEmail(priority: $priority) {
            email {
              from
              subject
              ... @include(if: $asyncResolver) {
                asyncSubject
              }
            }
            ... @defer(if: $shouldDefer) {
              inbox {
                emails @include(if: $shouldStream) @stream(if: $shouldStream)
                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, variable_values=variable_values)


DummyQueryType = GraphQLObjectType("Query", {"dummy": GraphQLField(GraphQLString)})


def subscribe_with_bad_fn(
    subscribe_fn: Callable,
) -> AwaitableOrValue[ExecutionResult | AsyncIterable[Any]]:
    schema = GraphQLSchema(
        query=DummyQueryType,
        subscription=GraphQLObjectType(
            "Subscription",
            {"foo": GraphQLField(GraphQLString, subscribe=subscribe_fn)},
        ),
    )
    document = parse("subscription { foo }")
    return subscribe_with_bad_args(schema, document)


def subscribe_with_bad_args(
    schema: GraphQLSchema,
    document: DocumentNode,
    variable_values: dict[str, Any] | None = None,
):
    executor = Executor.build(schema, document, raw_variable_values=variable_values)
    source_event_stream_result = (
        ExecutionResult(None, errors=executor)
        if isinstance(executor, list)
        else create_source_event_stream(executor)
    )

    return assert_equal_awaitables_or_values(
        subscribe(schema, document, variable_values=variable_values),
        source_event_stream_result,
    )


# Check all error cases when initializing the subscription.
def describe_subscription_initialization_phase():
    def throws_for_legacy_execution_args_passed_to_create_source_event_stream():
        schema = GraphQLSchema(
            query=DummyQueryType,
            subscription=GraphQLObjectType(
                "Subscription",
                {"foo": GraphQLField(GraphQLString)},
            ),
        )

        with pytest.raises(GraphQLError) as exc_info:
            create_source_event_stream(
                {  # type: ignore[arg-type]
                    "schema": schema,
                    "document": parse("subscription { foo }"),
                }
            )

        assert str(exc_info.value) == (
            "Passing execution arguments to create_source_event_stream()"
            " was removed in graphql-core version 3.3;"
            " call Executor.build() first and pass the result instead,"
            " or use subscribe() for the full subscription pipeline."
        )

    def throws_when_subscribe_is_called_with_a_non_subscription_operation():
        schema = GraphQLSchema(query=DummyQueryType)

        with pytest.raises(GraphQLError) as exc_info:
            subscribe(schema, parse("{ dummy }"))

        assert str(exc_info.value) == "Expected subscription operation."

    async def accepts_positional_arguments():
        document = parse(
            """
            subscription {
              importantEmail
            }
            """
        )

        async def empty_async_iterable(_info):
            for value in ():  # type: ignore
                yield value  # pragma: no cover

        ai = subscribe(email_schema, document, {"importantEmail": empty_async_iterable})
        assert isinstance(ai, AsyncIterator)

        with pytest.raises(StopAsyncIteration):
            await anext(ai)
        await ai.aclose()  # type: ignore

    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 = subscribe(
            schema, parse("subscription { foo }"), {"foo": foo_generator}
        )
        assert isinstance(subscription, AsyncIterator)

        assert await anext(subscription) == ({"foo": "FooValue"}, None)

        await subscription.aclose()  # type: ignore

    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 = subscribe(schema, parse("subscription { foo }"))
        assert isinstance(subscription, AsyncIterator)

        assert await anext(subscription) == ({"foo": "FooValue"}, None)

        await subscription.aclose()  # type: ignore

    async def accepts_type_definition_with_async_subscribe_function():
        async def foo_generator(_obj, _info):
            await asyncio.sleep(0)
            yield {"foo": "FooValue"}

        async def subscribe_fn(obj, info):
            await asyncio.sleep(0)
            return foo_generator(obj, info)

        schema = GraphQLSchema(
            query=DummyQueryType,
            subscription=GraphQLObjectType(
                "Subscription",
                {"foo": GraphQLField(GraphQLString, subscribe=subscribe_fn)},
            ),
        )

        awaitable = subscribe(schema, parse("subscription { foo }"))
        assert is_awaitable(awaitable)

        subscription = await awaitable
        assert isinstance(subscription, AsyncIterator)

        assert await anext(subscription) == ({"foo": "FooValue"}, None)

        await subscription.aclose()  # type: ignore

    async def maps_a_source_stream_with_a_custom_root_selection_set_executor():
        schema = GraphQLSchema(
            query=DummyQueryType,
            subscription=GraphQLObjectType(
                "Subscription", {"foo": GraphQLField(GraphQLString)}
            ),
        )

        async def foo_generator(_info):
            yield {"foo": "FooValue"}

        count = 0

        def root_selection_set_executor(context):
            nonlocal count
            count += 1
            return execute_subscription_event(context)

        executor = Executor.build(
            schema,
            parse("subscription { foo }"),
            {"foo": foo_generator},
        )
        assert isinstance(executor, Executor)

        result_or_stream = create_source_event_stream(executor)
        assert isinstance(result_or_stream, AsyncIterable)

        subscription = map_source_to_response_event(
            executor, result_or_stream, root_selection_set_executor
        )
        assert isinstance(subscription, AsyncIterator)

        assert await anext(subscription) == ({"foo": "FooValue"}, None)

        with pytest.raises(StopAsyncIteration):
            await anext(subscription)

        assert count == 1

    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 = subscribe(schema, parse("subscription { foo bar }"))
        assert isinstance(subscription, AsyncIterator)

        assert await anext(subscription) == (
            {"foo": "FooValue", "bar": None},
            None,
        )

        assert did_resolve == {"foo": True, "bar": False}

        await subscription.aclose()  # type: ignore

    async def resolves_to_an_error_if_schema_does_not_support_subscriptions():
        schema = GraphQLSchema(query=DummyQueryType)
        document = parse("subscription { unknownField }")

        result = subscribe_with_bad_args(schema, document)

        assert result == (
            None,
            [
                {
                    "message": "Schema is not configured to execute"
                    " subscription operation.",
                    "locations": [(1, 1)],
                }
            ],
        )

    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 = subscribe_with_bad_args(schema, document)
        assert result == (
            None,
            [
                {
                    "message": "The subscription field 'unknownField' is not defined.",
                    "locations": [(1, 16)],
                }
            ],
        )

    async def should_pass_through_unexpected_errors_thrown_in_subscribe():
        schema = GraphQLSchema(
            query=DummyQueryType,
            subscription=GraphQLObjectType(
                "Subscription", {"foo": GraphQLField(GraphQLString)}
            ),
        )
        with pytest.raises(AttributeError):
            subscribe_with_bad_args(schema=schema, document={})  # type: ignore

    async def throws_an_error_if_subscribe_does_not_return_an_iterator():
        expected_result = (
            None,
            [
                {
                    "message": "Subscription field must return AsyncIterable."
                    " Received: 'test'.",
                    "locations": [(1, 16)],
                    "path": ["foo"],
                }
            ],
        )

        def sync_fn(_obj, _info):
            return "test"

        assert subscribe_with_bad_fn(sync_fn) == expected_result

        async def async_fn(obj, info):
            return sync_fn(obj, info)

        result = subscribe_with_bad_fn(async_fn)
        assert is_awaitable(result)
        assert await result == expected_result

        del result
        cleanup()

    async def resolves_to_an_error_for_subscription_resolver_errors():
        expected_result = (
            None,
            [
                {
                    "message": "test error",
                    "locations": [(1, 16)],
                    "path": ["foo"],
                }
            ],
        )

        # Returning an error
        def return_error(*_args):
            return TypeError("test error")

        assert subscribe_with_bad_fn(return_error) == expected_result

        # Throwing an error
        def throw_error(*_args):
            raise TypeError("test error")

        assert subscribe_with_bad_fn(throw_error) == expected_result

        # Resolving to an error
        async def resolve_to_error(*args):
            return return_error(*args)

        result = subscribe_with_bad_fn(resolve_to_error)
        assert is_awaitable(result)
        assert await result == expected_result

        # Rejecting with an error

        async def reject_with_error(*args):
            return throw_error(*args)

        result = subscribe_with_bad_fn(reject_with_error)
        assert is_awaitable(result)
        assert await result == expected_result

    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 = subscribe_with_bad_args(
            schema, document, variable_values=variable_values
        )

        assert result == (
            None,
            [
                {
                    "message": "Variable '$arg' has invalid value:"
                    " Int cannot represent non-integer value: 'meow'",
                    "locations": [(2, 27)],
                }
            ],
        )

        assert result.errors[0].original_error


# Once a subscription returns a valid AsyncIterator, it can still yield errors.
def describe_subscription_publish_phase():
    async def produces_a_payload_for_multiple_subscribe_in_same_subscription():
        pubsub = SimplePubSub()

        subscription = create_subscription(pubsub)
        assert isinstance(subscription, AsyncIterator)

        second_subscription = create_subscription(pubsub)
        assert isinstance(second_subscription, AsyncIterator)

        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)

    async def produces_a_payload_when_queried_fields_are_async():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub, {"asyncResolver": True})
        assert isinstance(subscription, AsyncIterator)

        assert (
            pubsub.emit(
                {
                    "from": "yuzhi@graphql.org",
                    "subject": "Alright",
                    "message": "Tests are good",
                    "unread": True,
                }
            )
            is True
        )

        assert await anext(subscription) == (
            {
                "importantEmail": {
                    "email": {
                        "from": "yuzhi@graphql.org",
                        "subject": "Alright",
                        "asyncSubject": "Alright",
                    },
                    "inbox": {"unread": 1, "total": 2},
                }
            },
            None,
        )

        await subscription.aclose()  # type: ignore
        with pytest.raises(StopAsyncIteration):
            await anext(subscription)

    async def produces_a_payload_per_subscription_event():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub)
        assert isinstance(subscription, AsyncIterator)

        # 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.
        await subscription.aclose()  # type: ignore

        # 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 pytest.raises(StopAsyncIteration):
            assert await anext(subscription)

    async def subscribe_function_returns_errors_with_defer():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub, {"shouldDefer": True})
        assert isinstance(subscription, AsyncIterator)

        # 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
        )

        error_result = (
            {"importantEmail": None},
            [
                {
                    "message": "`@defer` directive not supported"
                    " on subscription operations."
                    " Disable `@defer` by setting the `if` argument to `false`.",
                    "locations": [(8, 11)],
                    "path": ["importantEmail"],
                }
            ],
        )

        # The previously waited on payload now has a value.
        result = await payload
        assert result == error_result

        # Another new email arrives,
        # after all incrementally delivered payloads are received.
        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.
        result = await anext(subscription)
        assert result == error_result

        await subscription.aclose()  # type: ignore

        # Awaiting a subscription after closing it results in completed results.
        with pytest.raises(StopAsyncIteration):
            assert await anext(subscription)

    async def subscribe_function_returns_errors_with_stream():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub, {"shouldStream": True})
        assert isinstance(subscription, AsyncIterator)

        # 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": {"emails": None, "unread": 1, "total": 2},
                }
            },
            [
                {
                    "message": "`@stream` directive not supported"
                    " on subscription operations."
                    " Disable `@stream` by setting the `if` argument to `false`.",
                    "locations": [(18, 17)],
                    "path": ["importantEmail", "inbox", "emails"],
                }
            ],
        )

        # Another new email arrives,
        # after all incrementally delivered payloads are received.
        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": {"emails": None, "unread": 2, "total": 3},
                }
            },
            [
                {
                    "message": "`@stream` directive not supported"
                    " on subscription operations."
                    " Disable `@stream` by setting the `if` argument to `false`.",
                    "locations": [(18, 17)],
                    "path": ["importantEmail", "inbox", "emails"],
                }
            ],
        )

        # The client disconnects before the deferred payload is consumed.
        await subscription.aclose()  # type: ignore

        # Awaiting a subscription after closing it results in completed results.
        with pytest.raises(StopAsyncIteration):
            assert await anext(subscription)

    async def produces_a_payload_when_there_are_multiple_events():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub)
        assert isinstance(subscription, AsyncIterator)

        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,
        )

    async def should_not_trigger_when_subscription_is_already_done():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub)
        assert isinstance(subscription, AsyncIterator)

        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()  # type: ignore

        # A new email arrives!
        assert (
            pubsub.emit(
                {
                    "from": "yuzhi@graphql.org",
                    "subject": "Alright 2",
                    "message": "Tests are good 2",
                    "unread": True,
                }
            )
            is False
        )

        with pytest.raises(StopAsyncIteration):
            await payload

    async def should_not_trigger_when_subscription_is_thrown():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub)
        assert isinstance(subscription, AsyncIterator)

        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 pytest.raises(RuntimeError) as exc_info:
            await subscription.athrow(RuntimeError("ouch"))  # type: ignore
        assert str(exc_info.value) == "ouch"

        with pytest.raises(StopAsyncIteration):
            await payload

    async def event_order_is_correct_for_multiple_publishes():
        pubsub = SimplePubSub()
        subscription = create_subscription(pubsub)
        assert isinstance(subscription, AsyncIterator)

        payload = anext(subscription)

        # A new email arrives!
        assert (
            pubsub.emit(
                {
                    "from": "yuzhi@graphql.org",
                    "subject": "Message",
                    "message": "Tests are good",
                    "unread": True,
                }
            )