Fix async tests

Author: SmileyChris

graphql_ws/base.py

@@ -23,6 +23,8 @@ class ConnectionClosedException(Exception):
 
 
 class BaseConnectionContext(object):
+    transport_ws_protocol = False
+
     def __init__(self, ws, request_context=None):
         self.ws = ws
         self.operations = {}

setup.cfg

@@ -50,9 +50,8 @@ test =
     pytest-asyncio; python_version>="3.4"
     graphene>=2.0,<3
     gevent
-    graphene>=2.0
     graphene_django
-    mock; python_version<"3"
+    mock; python_version<"3.8"
     django==1.11.*; python_version<"3"
     channels==1.*; python_version<"3"
     django==2.*; python_version>="3"

tests/test_base_async.py

@@ -10,9 +10,10 @@
 pytestmark = pytest.mark.asyncio
 
 
-class AsyncMock(mock.MagicMock):
-    async def __call__(self, *args, **kwargs):
-        return super().__call__(*args, **kwargs)
+try:
+    from unittest.mock import AsyncMock  # Python 3.8+
+except ImportError:
+    from mock import AsyncMock
 
 
 class TstServer(base_async.BaseAsyncSubscriptionServer):
@@ -26,75 +27,78 @@ def server():
 
 
 async def test_terminate(server: TstServer):
-    context = AsyncMock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     await server.on_connection_terminate(connection_context=context, op_id=1)
     context.close.assert_called_with(1011)
 
 
 async def test_send_error(server: TstServer):
-    context = AsyncMock()
-    context.has_operation = mock.Mock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     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"}}
     )
 
 
-async def test_message(server):
+async def test_message(server: TstServer):
     server.process_message = AsyncMock()
-    context = AsyncMock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     msg = {"id": 1, "type": base.GQL_CONNECTION_INIT, "payload": ""}
     await server.on_message(context, msg)
     server.process_message.assert_called_with(context, msg)
 
 
-async def test_message_str(server):
+async def test_message_str(server: TstServer):
     server.process_message = AsyncMock()
-    context = AsyncMock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     msg = {"id": 1, "type": base.GQL_CONNECTION_INIT, "payload": ""}
     await server.on_message(context, json.dumps(msg))
     server.process_message.assert_called_with(context, msg)
 
 
-async def test_message_invalid(server):
+async def test_message_invalid(server: TstServer):
     server.send_error = AsyncMock()
-    await server.on_message(connection_context=None, message="'not-json")
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
+    await server.on_message(context, message="'not-json")
     assert server.send_error.called
 
 
-async def test_resolver(server):
+async def test_resolver(server: TstServer):
     server.send_message = AsyncMock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     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
+        context, op_id=1, execution_result=result
     )
     assert server.send_message.called
 
 
 @pytest.mark.asyncio
-async def test_resolver_with_promise(server):
+async def test_resolver_with_promise(server: TstServer):
     server.send_message = AsyncMock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     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
+        context, 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):
+async def test_resolver_with_nested_promise(server: TstServer):
     server.send_message = AsyncMock()
+    context = AsyncMock(spec=base_async.BaseAsyncConnectionContext)
     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
+        context, op_id=1, execution_result=result
     )
     assert server.send_message.called
     assert result.data == {'test': [1, {'in': 2}]}

tests/test_graphql_ws.py

@@ -165,7 +165,7 @@ def test_build_message_partial(ss):
         ss.build_message(id=None, op_type=None, payload=None)
 
 
-def test_send_execution_result(ss):
+def test_send_execution_result(ss, cc):
     ss.execution_result_to_dict = mock.Mock()
     ss.execution_result_to_dict.return_value = {"res": "ult"}
     ss.send_message = mock.Mock()