Allow to pass arguments to websockets.connect + some ssl tests + some cleaning (#83)

Author: leszekhanusz

CONTRIBUTING.md

@@ -12,10 +12,10 @@ We welcome all kinds of contributions:
 
 ## Getting started
 
-If you have a specific contribution in mind, be sure to check the 
-[issues](https://github.com/graphql-python/gql/issues) 
-and [pull requests](https://github.com/graphql-python/gql/pulls) 
-in progress - someone could already be working on something similar 
+If you have a specific contribution in mind, be sure to check the
+[issues](https://github.com/graphql-python/gql/issues)
+and [pull requests](https://github.com/graphql-python/gql/pulls)
+in progress - someone could already be working on something similar
 and you can help out.
 
 ## Project setup
@@ -31,10 +31,10 @@ virtualenv gql-dev
 Activate the virtualenv and install dependencies by running:
 
 ```console
-python pip install -e[dev]
+python pip install -e.[dev]
 ```
 
-If you are using Linux or MacOS, you can make use of Makefile command 
+If you are using Linux or MacOS, you can make use of Makefile command
 `make dev-setup`, which is a shortcut for the above python command.
 
 ### Development on Conda
@@ -55,7 +55,7 @@ pip install -e.[dev]
 
 And you ready to start development!
 
-<!-- TODO: Provide environment.yml file for conda env --> 
+<!-- TODO: Provide environment.yml file for conda env -->
 
 ## Running tests
 
@@ -65,7 +65,7 @@ After developing, the full test suite can be evaluated by running:
 pytest tests --cov=gql -vv
 ```
 
-If you are using Linux or MacOS, you can make use of Makefile command 
+If you are using Linux or MacOS, you can make use of Makefile command
 `make tests`, which is a shortcut for the above python command.
 
 You can also test on several python environments by using tox.
@@ -77,8 +77,8 @@ Install tox:
 pip install tox
 ```
 
-Run `tox` on your virtualenv (do not forget to activate it!) 
-and that's it! 
+Run `tox` on your virtualenv (do not forget to activate it!)
+and that's it!
 
 ### Running tox on Conda
 
@@ -93,5 +93,5 @@ This install tox underneath so no need to install it before.
 
 Then uncomment the `requires = tox-conda` line on `tox.ini` file.
 
-Run `tox` and you will see all the environments being created 
-and all passing tests. :rocket:
+Run `tox` and you will see all the environments being created
+and all passing tests. :rocket:

MANIFEST.in

@@ -13,7 +13,7 @@ include tox.ini
 include scripts/gql-cli
 
 recursive-include tests *.py *.yaml *.graphql
-recursive-include tests_py36 *.py
+recursive-include tests_py36 *.py *.cnf *.pem
 
 prune gql-checker
 

Makefile

@@ -1,8 +1,13 @@
+.PHONY: clean tests
+
 dev-setup:
 	python pip install -e ".[test]"
 
 tests:
-	pytest tests --cov=gql -vv
+	pytest tests tests_py36 --cov=gql --cov-report=term-missing -vv
+
+all_tests:
+	pytest tests tests_py36 --cov=gql --cov-report=term-missing --run-online -vv
 
 clean:
 	find . -name "*.pyc" -delete

README.md

@@ -67,10 +67,6 @@ from gql.transport.requests import RequestsHTTPTransport
 
 sample_transport=RequestsHTTPTransport(
     url='https://countries.trevorblades.com/',
-    use_json=True,
-    headers={
-        "Content-type": "application/json",
-    },
     verify=False,
     retries=3,
 )
@@ -215,7 +211,6 @@ async def main():
 
     sample_transport = WebsocketsTransport(
         url='wss://countries.trevorblades.com/graphql',
-        ssl=True,
         headers={'Authorization': 'token'}
     )
 
@@ -262,8 +257,7 @@ import ssl
 
 sample_transport = WebsocketsTransport(
     url='wss://SERVER_URL:SERVER_PORT/graphql',
-    headers={'Authorization': 'token'},
-    ssl=True
+    headers={'Authorization': 'token'}
 )
 ```
 
@@ -298,8 +292,7 @@ There are two ways to send authentication tokens with websockets depending on th
 ```python
 sample_transport = WebsocketsTransport(
     url='wss://SERVER_URL:SERVER_PORT/graphql',
-    headers={'Authorization': 'token'},
-    ssl=True
+    headers={'Authorization': 'token'}
 )
 ```
 
@@ -308,8 +301,7 @@ sample_transport = WebsocketsTransport(
 ```python
 sample_transport = WebsocketsTransport(
     url='wss://SERVER_URL:SERVER_PORT/graphql',
-    init_payload={'Authorization': 'token'},
-    ssl=True
+    init_payload={'Authorization': 'token'}
 )
 ```
 

gql/transport/aiohttp.py

@@ -35,7 +35,7 @@ def __init__(
         auth: Optional[BasicAuth] = None,
         ssl: Union[SSLContext, bool, Fingerprint] = False,
         timeout: Optional[int] = None,
-        **kwargs,
+        client_session_args: Dict[str, Any] = {},
     ) -> None:
         """Initialize the transport with the given aiohttp parameters.
 
@@ -44,15 +44,15 @@ def __init__(
         :param cookies: Dict of HTTP cookies.
         :param auth: BasicAuth object to enable Basic HTTP auth if needed
         :param ssl: ssl_context of the connection. Use ssl=False to disable encryption
-        :param kwargs: Other parameters forwarded to aiohttp.ClientSession
+        :param client_session_args: Dict of extra parameters passed to aiohttp.ClientSession
         """
         self.url: str = url
         self.headers: Optional[LooseHeaders] = headers
         self.cookies: Optional[LooseCookies] = cookies
         self.auth: Optional[BasicAuth] = auth
         self.ssl: Union[SSLContext, bool, Fingerprint] = ssl
         self.timeout: Optional[int] = timeout
-        self.kwargs = kwargs
+        self.client_session_args = client_session_args
 
         self.session: Optional[aiohttp.ClientSession] = None
 
@@ -78,7 +78,7 @@ async def connect(self) -> None:
                 )
 
             # Adding custom parameters passed from init
-            client_session_args.update(self.kwargs)
+            client_session_args.update(self.client_session_args)
 
             self.session = aiohttp.ClientSession(**client_session_args)
 
@@ -95,7 +95,7 @@ async def execute(
         document: Document,
         variable_values: Optional[Dict[str, str]] = None,
         operation_name: Optional[str] = None,
-        **kwargs,
+        extra_args: Dict[str, Any] = {},
     ) -> ExecutionResult:
         """Execute the provided document AST against the configured remote server.
         This uses the aiohttp library to perform a HTTP POST request asynchronously to the remote server.
@@ -114,8 +114,8 @@ async def execute(
             "json": payload,
         }
 
-        # Pass kwargs to aiohttp post method
-        post_args.update(kwargs)
+        # Pass post_args to aiohttp post method
+        post_args.update(extra_args)
 
         if self.session is None:
             raise TransportClosed("Transport is not connected")

gql/transport/websockets.py

@@ -97,6 +97,7 @@ def __init__(
         connect_timeout: int = 10,
         close_timeout: int = 10,
         ack_timeout: int = 10,
+        connect_args: Dict[str, Any] = {},
     ) -> None:
         """Initialize the transport with the given request parameters.
 
@@ -107,6 +108,7 @@ def __init__(
         :param connect_timeout: Timeout in seconds for the establishment of the websocket connection.
         :param close_timeout: Timeout in seconds for the close.
         :param ack_timeout: Timeout in seconds to wait for the connection_ack message from the server.
+        :param connect_args: Other parameters forwarded to websockets.connect
         """
         self.url: str = url
         self.ssl: Union[SSLContext, bool] = ssl
@@ -117,6 +119,8 @@ def __init__(
         self.close_timeout: int = close_timeout
         self.ack_timeout: int = ack_timeout
 
+        self.connect_args = connect_args
+
         self.websocket: Optional[WebSocketClientProtocol] = None
         self.next_query_id: int = 1
         self.listeners: Dict[int, ListenerQueue] = {}
@@ -460,16 +464,27 @@ async def connect(self) -> None:
 
         if self.websocket is None:
 
+            # If the ssl parameter is not provided, generate the ssl value depending on the url
+            ssl: Optional[Union[SSLContext, bool]]
+            if self.ssl:
+                ssl = self.ssl
+            else:
+                ssl = True if self.url.startswith("wss") else None
+
+            # Set default arguments used in the websockets.connect call
+            connect_args: Dict[str, Any] = {
+                "ssl": ssl,
+                "extra_headers": self.headers,
+                "subprotocols": [GRAPHQLWS_SUBPROTOCOL],
+            }
+
+            # Adding custom parameters passed from init
+            connect_args.update(self.connect_args)
+
             # Connection to the specified url
             # Generate a TimeoutError if taking more than connect_timeout seconds
             self.websocket = await asyncio.wait_for(
-                websockets.connect(
-                    self.url,
-                    ssl=self.ssl if self.ssl else None,
-                    extra_headers=self.headers,
-                    subprotocols=[GRAPHQLWS_SUBPROTOCOL],
-                ),
-                self.connect_timeout,
+                websockets.connect(self.url, **connect_args,), self.connect_timeout,
             )
 
             self.next_query_id = 1

tests_py36/conftest.py

@@ -2,6 +2,8 @@
 import json
 import logging
 import os
+import pathlib
+import ssl
 import types
 
 import pytest
@@ -79,12 +81,35 @@ class TestServer:
     Will allow us to test our client by simulating different correct and incorrect server responses
     """
 
+    def __init__(self, with_ssl: bool = False):
+        self.with_ssl = with_ssl
+
     async def start(self, handler):
 
         print("Starting server")
 
+        extra_serve_args = {}
+
+        if self.with_ssl:
+            # This is a copy of certificate from websockets tests folder
+            #
+            # Generate TLS certificate with:
+            # $ openssl req -x509 -config test_localhost.cnf -days 15340 -newkey rsa:2048 \
+            #       -out test_localhost.crt -keyout test_localhost.key
+            # $ cat test_localhost.key test_localhost.crt > test_localhost.pem
+            # $ rm test_localhost.key test_localhost.crt
+            self.testcert = bytes(
+                pathlib.Path(__file__).with_name("test_localhost.pem")
+            )
+            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+            ssl_context.load_cert_chain(self.testcert)
+
+            extra_serve_args["ssl"] = ssl_context
+
         # Start a server with a random open port
-        self.start_server = websockets.server.serve(handler, "localhost", 0)
+        self.start_server = websockets.server.serve(
+            handler, "localhost", 0, **extra_serve_args
+        )
 
         # Wait that the server is started
         self.server = await self.start_server
@@ -137,13 +162,9 @@ async def wait_connection_terminate(ws):
         assert json_result["type"] == "connection_terminate"
 
 
-@pytest.fixture
-async def server(request):
-    """server is a fixture used to start a dummy server to test the client behaviour.
-
-    It can take as argument either a handler function for the websocket server for complete control
-    OR an array of answers to be sent by the default server handler
-    """
+def get_server_handler(request):
+    """ Get the server handler provided from test or use the default
+    server handler if the test provides only an array of answers"""
 
     if isinstance(request.param, types.FunctionType):
         server_handler = request.param
@@ -179,6 +200,42 @@ async def default_server_handler(ws, path):
 
         server_handler = default_server_handler
 
+    return server_handler
+
+
+@pytest.fixture
+async def ws_ssl_server(request):
+    """websockets server fixture using ssl
+
+    It can take as argument either a handler function for the websocket server for complete control
+    OR an array of answers to be sent by the default server handler
+    """
+
+    server_handler = get_server_handler(request)
+
+    try:
+        test_server = TestServer(with_ssl=True)
+
+        # Starting the server with the fixture param as the handler function
+        await test_server.start(server_handler)
+
+        yield test_server
+    except Exception as e:
+        print("Exception received in server fixture: " + str(e))
+    finally:
+        await test_server.stop()
+
+
+@pytest.fixture
+async def server(request):
+    """server is a fixture used to start a dummy server to test the client behaviour.
+
+    It can take as argument either a handler function for the websocket server for complete control
+    OR an array of answers to be sent by the default server handler
+    """
+
+    server_handler = get_server_handler(request)
+
     try:
         test_server = TestServer()