Add tests

Author: leszekhanusz

tests/test_aiohttp.py

@@ -948,3 +948,35 @@ async def handler(request):
     expected_error = "Syntax Error: Unexpected Name 'BLAHBLAH'"
 
     assert expected_error in captured_err
+
+
+query1_server_answer_with_extensions = (
+    f'{{"data":{query1_server_answer_data}, "extensions":{{"key1": "val1"}}}}'
+)
+
+
+@pytest.mark.asyncio
+async def test_aiohttp_query_with_extensions(event_loop, aiohttp_server):
+    from aiohttp import web
+    from gql.transport.aiohttp import AIOHTTPTransport
+
+    async def handler(request):
+        return web.Response(
+            text=query1_server_answer_with_extensions, content_type="application/json"
+        )
+
+    app = web.Application()
+    app.router.add_route("POST", "/", handler)
+    server = await aiohttp_server(app)
+
+    url = server.make_url("/")
+
+    sample_transport = AIOHTTPTransport(url=url, timeout=10)
+
+    async with Client(transport=sample_transport,) as session:
+
+        query = gql(query1_str)
+
+        execution_result = await session._execute(query)
+
+        assert execution_result.extensions["key1"] == "val1"

tests/test_requests.py

@@ -216,3 +216,47 @@ def test_code():
             sample_transport.execute(query)
 
     await run_sync_test(event_loop, server, test_code)
+
+
+query1_server_answer_with_extensions = (
+    '{"data":{"continents":['
+    '{"code":"AF","name":"Africa"},{"code":"AN","name":"Antarctica"},'
+    '{"code":"AS","name":"Asia"},{"code":"EU","name":"Europe"},'
+    '{"code":"NA","name":"North America"},{"code":"OC","name":"Oceania"},'
+    '{"code":"SA","name":"South America"}]},'
+    '"extensions": {"key1": "val1"}'
+    "}"
+)
+
+
+@pytest.mark.aiohttp
+@pytest.mark.asyncio
+async def test_requests_query_with_extensions(
+    event_loop, aiohttp_server, run_sync_test
+):
+    from aiohttp import web
+    from gql.transport.requests import RequestsHTTPTransport
+
+    async def handler(request):
+        return web.Response(
+            text=query1_server_answer_with_extensions, content_type="application/json"
+        )
+
+    app = web.Application()
+    app.router.add_route("POST", "/", handler)
+    server = await aiohttp_server(app)
+
+    url = server.make_url("/")
+
+    def test_code():
+        sample_transport = RequestsHTTPTransport(url=url)
+
+        with Client(transport=sample_transport,) as session:
+
+            query = gql(query1_str)
+
+            execution_result = session._execute(query)
+
+            assert execution_result.extensions["key1"] == "val1"
+
+    await run_sync_test(event_loop, server, test_code)

tests/test_websocket_query.py

@@ -569,3 +569,33 @@ async def test_websocket_using_cli(event_loop, server, monkeypatch, capsys):
     received_answer = json.loads(captured_out)
 
     assert received_answer == expected_answer
+
+
+query1_server_answer_with_extensions = (
+    '{{"type":"data","id":"{query_id}","payload":{{"data":{{"continents":['
+    '{{"code":"AF","name":"Africa"}},{{"code":"AN","name":"Antarctica"}},'
+    '{{"code":"AS","name":"Asia"}},{{"code":"EU","name":"Europe"}},'
+    '{{"code":"NA","name":"North America"}},{{"code":"OC","name":"Oceania"}},'
+    '{{"code":"SA","name":"South America"}}]}},'
+    '"extensions": {{"key1": "val1"}}}}}}'
+)
+
+server1_answers_with_extensions = [
+    query1_server_answer_with_extensions,
+]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("server", [server1_answers_with_extensions], indirect=True)
+@pytest.mark.parametrize("query_str", [query1_str])
+async def test_websocket_simple_query_with_extensions(
+    event_loop, client_and_server, query_str
+):
+
+    session, server = client_and_server
+
+    query = gql(query_str)
+
+    execution_result = await session._execute(query)
+
+    assert execution_result.extensions["key1"] == "val1"