added docs and tests for websocket_kwargs

Author: voith

docs/providers.rst

@@ -224,15 +224,30 @@ IPCProvider
 WebsocketProvider
 ~~~~~~~~~~~~~~~~~
 
-.. py:class:: web3.providers.websocket.WebsocketProvider(endpoint_uri)
+.. py:class:: web3.providers.websocket.WebsocketProvider(endpoint_uri[, websocket_kwargs])
 
     This provider handles interactions with an WS or WSS based JSON-RPC server.
 
+    * ``endpoint_uri`` should be the full URI to the RPC endpoint such as
+      ``'ws://localhost:8546'``.
+    * ``websocket_kwargs`` this should be a dictionary of keyword arguments which
+      will be passed onto the ws/wss websocket connection.
+
     .. code-block:: python
 
         >>> from web3 import Web3
         >>> web3 = Web3(Web3.WebsocketProvider("ws://127.0.0.1:8546")
 
+    Under the hood, the ``WebsocketProvider`` uses the python websockets library for
+    making requests.  If you would like to modify how requests are made, you can
+    use the ``websocket_kwargs`` to do so.  A common use case for this is increasing
+    the timeout for each request.
+
+
+    .. code-block:: python
+
+        >>> from web3 import Web3
+        >>> web3 = Web3(Web3.WebsocketProvider("http://127.0.0.1:8546", websocket_kwargs={'timeout': 60}))
 
 .. py:currentmodule:: web3.providers.eth_tester
 

tests/core/providers/test_websocket_provider.py

@@ -0,0 +1,15 @@
+import pytest
+
+from web3.exceptions import (
+    ValidationError,
+)
+from web3.providers.websocket import (
+    WebsocketProvider,
+)
+
+
+def test_restricted_websocket_kwargs():
+    invalid_kwargs = {'uri': 'ws://127.0.0.1:8546'}
+    re_exc_message = r'.*found: {0}*'.format(set(invalid_kwargs.keys()))
+    with pytest.raises(ValidationError, match=re_exc_message):
+        WebsocketProvider(websocket_kwargs=invalid_kwargs)

tests/integration/common.py

@@ -0,0 +1,17 @@
+import pytest
+
+from websockets.exceptions import (
+    ConnectionClosed,
+)
+
+from web3 import Web3
+
+
+class MiscWebsocketTests:
+
+    def test_websocket_max_size_error(self, web3, endpoint_uri):
+        w3 = Web3(Web3.WebsocketProvider(
+            endpoint_uri=endpoint_uri, websocket_kwargs={'max_size': 1})
+        )
+        with pytest.raises(ConnectionClosed):
+            w3.eth.getBlock(0)

tests/integration/go_ethereum/test_goethereum_ws.py

@@ -1,5 +1,8 @@
 import pytest
 
+from tests.integration.common import (
+    MiscWebsocketTests,
+)
 from tests.integration.utils import (
     wait_for_ws,
 )
@@ -65,3 +68,7 @@ class TestGoEthereumNetModuleTest(GoEthereumNetModuleTest):
 
 class TestGoEthereumPersonalModuleTest(GoEthereumPersonalModuleTest):
     pass
+
+
+class TestMiscWebsocketTests(MiscWebsocketTests):
+    pass

tests/integration/parity/test_parity_ws.py

@@ -1,6 +1,9 @@
 import os
 import pytest
 
+from tests.integration.common import (
+    MiscWebsocketTests,
+)
 from tests.integration.utils import (
     wait_for_ws,
 )
@@ -96,3 +99,7 @@ class TestParityPersonalModuleTest(ParityPersonalModuleTest):
 
 class TestParityTraceModuleTest(ParityTraceModuleTest):
     pass
+
+
+class TestMiscWebsocketTests(MiscWebsocketTests):
+    pass

web3/providers/websocket.py

@@ -8,10 +8,15 @@
 
 import websockets
 
+from web3.exceptions import (
+    ValidationError,
+)
 from web3.providers.base import (
     JSONBaseProvider,
 )
 
+RESTRICTED_WEBSOCKET_KWARGS = {'uri', 'loop'}
+
 
 def _start_event_loop(loop):
     asyncio.set_event_loop(loop)
@@ -32,15 +37,17 @@ def get_default_endpoint():
 
 class PersistentWebSocket:
 
-    def __init__(self, endpoint_uri, loop, **kwargs):
+    def __init__(self, endpoint_uri, loop, websocket_kwargs):
         self.ws = None
         self.endpoint_uri = endpoint_uri
         self.loop = loop
-        self.kwargs = kwargs
+        self.websocket_kwargs = websocket_kwargs
 
     async def __aenter__(self):
         if self.ws is None:
-            self.ws = await websockets.connect(uri=self.endpoint_uri, loop=self.loop, **self.kwargs)
+            self.ws = await websockets.connect(
+                uri=self.endpoint_uri, loop=self.loop, **self.websocket_kwargs
+            )
         return self.ws
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
@@ -56,13 +63,26 @@ class WebsocketProvider(JSONBaseProvider):
     logger = logging.getLogger("web3.providers.WebsocketProvider")
     _loop = None
 
-    def __init__(self, endpoint_uri=None, **kwargs):
+    def __init__(self, endpoint_uri=None, websocket_kwargs=None):
         self.endpoint_uri = endpoint_uri
         if self.endpoint_uri is None:
             self.endpoint_uri = get_default_endpoint()
         if WebsocketProvider._loop is None:
             WebsocketProvider._loop = _get_threaded_loop()
-        self.conn = PersistentWebSocket(self.endpoint_uri, WebsocketProvider._loop, **kwargs)
+        if websocket_kwargs is None:
+            websocket_kwargs = {}
+        else:
+            found_restricted_keys = set(websocket_kwargs.keys()).intersection(
+                RESTRICTED_WEBSOCKET_KWARGS
+            )
+            if found_restricted_keys:
+                raise ValidationError(
+                    '{0} are not allowed in websocket_kwargs, '
+                    'found: {1}'.format(RESTRICTED_WEBSOCKET_KWARGS, found_restricted_keys)
+                )
+        self.conn = PersistentWebSocket(
+            self.endpoint_uri, WebsocketProvider._loop, websocket_kwargs
+        )
         super().__init__()
 
     def __str__(self):