Recursively resolve Promises, fix async tests

Author: SmileyChris

graphql_ws/base_async.py

@@ -1,16 +1,62 @@
 import asyncio
+import inspect
 from abc import ABC, abstractmethod
-from inspect import isawaitable
+from types import CoroutineType, GeneratorType
+from typing import Any, Union, List, Dict
 from weakref import WeakSet
 
 from graphql.execution.executors.asyncio import AsyncioExecutor
+from promise import Promise
 
 from graphql_ws import base
 
 from .constants import GQL_COMPLETE, GQL_CONNECTION_ACK, GQL_CONNECTION_ERROR
 from .observable_aiter import setup_observable_extension
 
 setup_observable_extension()
+CO_ITERABLE_COROUTINE = inspect.CO_ITERABLE_COROUTINE
+
+
+# Copied from graphql-core v3.1.0 (graphql/pyutils/is_awaitable.py)
+def is_awaitable(value: Any) -> bool:
+    """Return true if object can be passed to an ``await`` expression.
+    Instead of testing if the object is an instance of abc.Awaitable, it checks
+    the existence of an `__await__` attribute. This is much faster.
+    """
+    return (
+        # check for coroutine objects
+        isinstance(value, CoroutineType)
+        # check for old-style generator based coroutine objects
+        or isinstance(value, GeneratorType)
+        and bool(value.gi_code.co_flags & CO_ITERABLE_COROUTINE)
+        # check for other awaitables (e.g. futures)
+        or hasattr(value, "__await__")
+    )
+
+
+async def resolve(
+    data: Any, _container: Union[List, Dict] = None, _key: Union[str, int] = None
+) -> None:
+    """
+    Recursively wait on any awaitable children of a data element and resolve any
+    Promises.
+    """
+    if is_awaitable(data):
+        data = await data
+        if isinstance(data, Promise):
+            data = data.value  # type: Any
+        if _container is not None:
+            _container[_key] = data
+    if isinstance(data, dict):
+        items = data.items()
+    elif isinstance(data, list):
+        items = enumerate(data)
+    else:
+        items = None
+    if items is not None:
+        children = [resolve(child, _container=data, _key=key) for key, child in items]
+        if children:
+            await asyncio.wait(children)
 
 
 class BaseAsyncConnectionContext(base.BaseConnectionContext, ABC):
@@ -81,7 +127,7 @@ async def on_connection_init(self, connection_context, op_id, payload):
     async def on_start(self, connection_context, op_id, params):
         execution_result = self.execute(params)
 
-        if isawaitable(execution_result):
+        if is_awaitable(execution_result):
             execution_result = await execution_result
 
         if hasattr(execution_result, "__aiter__"):
@@ -120,3 +166,8 @@ async def on_stop(self, connection_context, op_id):
 
     async def on_operation_complete(self, connection_context, op_id):
         pass
+
+    async def send_execution_result(self, connection_context, op_id, execution_result):
+        # Resolve any pending promises
+        await resolve(execution_result.data)
+        await super().send_execution_result(connection_context, op_id, execution_result)

tests/test_base_async.py

@@ -1,59 +1,99 @@
 from unittest import mock
 
 import json
+import promise
 
 import pytest
 
-from graphql_ws import base
+from graphql_ws import base, base_async
 
+pytestmark = pytest.mark.asyncio
 
-def test_not_implemented():
-    server = base.BaseSubscriptionServer(schema=None)
-    with pytest.raises(NotImplementedError):
-        server.on_connection_init(connection_context=None, op_id=1, payload={})
-    with pytest.raises(NotImplementedError):
-        server.on_open(connection_context=None)
-    with pytest.raises(NotImplementedError):
-        server.on_stop(connection_context=None, op_id=1)
 
+class AsyncMock(mock.MagicMock):
+    async def __call__(self, *args, **kwargs):
+        return super().__call__(*args, **kwargs)
 
-def test_terminate():
-    server = base.BaseSubscriptionServer(schema=None)
 
-    context = mock.Mock()
-    server.on_connection_terminate(connection_context=context, op_id=1)
+class TestServer(base_async.BaseAsyncSubscriptionServer):
+    def handle(self, *args, **kwargs):
+        pass
+
+
+@pytest.fixture
+def server():
+    return TestServer(schema=None)
+
+
+async def test_terminate(server: TestServer):
+    context = AsyncMock()
+    await server.on_connection_terminate(connection_context=context, op_id=1)
     context.close.assert_called_with(1011)
 
 
-def test_send_error():
-    server = base.BaseSubscriptionServer(schema=None)
-    context = mock.Mock()
-    server.send_error(connection_context=context, op_id=1, error="test error")
+async def test_send_error(server: TestServer):
+    context = AsyncMock()
+    await server.send_error(connection_context=context, op_id=1, error="test error")
     context.send.assert_called_with(
         {"id": 1, "type": "error", "payload": {"message": "test error"}}
     )
 
 
-def test_message():
-    server = base.BaseSubscriptionServer(schema=None)
-    server.process_message = mock.Mock()
-    context = mock.Mock()
+async def test_message(server):
+    server.process_message = AsyncMock()
+    context = AsyncMock()
     msg = {"id": 1, "type": base.GQL_CONNECTION_INIT, "payload": ""}
-    server.on_message(context, msg)
+    await server.on_message(context, msg)
     server.process_message.assert_called_with(context, msg)
 
 
-def test_message_str():
-    server = base.BaseSubscriptionServer(schema=None)
-    server.process_message = mock.Mock()
-    context = mock.Mock()
+async def test_message_str(server):
+    server.process_message = AsyncMock()
+    context = AsyncMock()
     msg = {"id": 1, "type": base.GQL_CONNECTION_INIT, "payload": ""}
-    server.on_message(context, json.dumps(msg))
+    await server.on_message(context, json.dumps(msg))
     server.process_message.assert_called_with(context, msg)
 
 
-def test_message_invalid():
-    server = base.BaseSubscriptionServer(schema=None)
-    server.send_error = mock.Mock()
-    server.on_message(connection_context=None, message="'not-json")
+async def test_message_invalid(server):
+    server.send_error = AsyncMock()
+    await server.on_message(connection_context=None, message="'not-json")
     assert server.send_error.called
+
+
+async def test_resolver(server):
+    server.send_message = AsyncMock()
+    result = mock.Mock()
+    result.data = {"test": [1, 2]}
+    result.errors = None
+    await server.send_execution_result(
+        connection_context=None, op_id=1, execution_result=result
+    )
+    assert server.send_message.called
+
+
+@pytest.mark.asyncio
+async def test_resolver_with_promise(server):
+    server.send_message = AsyncMock()
+    result = mock.Mock()
+    result.data = {"test": [1, promise.Promise(lambda resolve, reject: resolve(2))]}
+    result.errors = None
+    await server.send_execution_result(
+        connection_context=None, op_id=1, execution_result=result
+    )
+    assert server.send_message.called
+    assert result.data == {'test': [1, 2]}
+
+
+async def test_resolver_with_nested_promise(server):
+    server.send_message = AsyncMock()
+    result = mock.Mock()
+    inner = promise.Promise(lambda resolve, reject: resolve(2))
+    outer = promise.Promise(lambda resolve, reject: resolve({'in': inner}))
+    result.data = {"test": [1, outer]}
+    result.errors = None
+    await server.send_execution_result(
+        connection_context=None, op_id=1, execution_result=result
+    )
+    assert server.send_message.called
+    assert result.data == {'test': [1, {'in': 2}]}