Merge pull request #1119 from voith/v4-add-websocket-timeout

carver · web-flow · commit e3e7510b4a75 · 2018-10-21T21:04:20.000-07:00
[BACKPORT TO V4] Add timeout for WebsocketProvider
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -8,7 +8,16 @@
     identity,
 )
 
+from .utils import (
+    get_open_port,
+)
+
 
 @pytest.fixture(scope="module", params=[lambda x: to_bytes(hexstr=x), identity])
 def address_conversion_func(request):
     return request.param
+
+
+@pytest.fixture()
+def open_port():
+    return get_open_port()
diff --git a/tests/core/providers/test_websocket_provider.py b/tests/core/providers/test_websocket_provider.py
@@ -1,5 +1,18 @@
+import asyncio
+from concurrent.futures import (
+    TimeoutError,
+)
 import pytest
+from threading import (
+    Thread,
+)
+
+import websockets
 
+from tests.utils import (
+    wait_for_ws,
+)
+from web3 import Web3
 from web3.exceptions import (
     ValidationError,
 )
@@ -8,6 +21,42 @@
 )
 
 
+@pytest.yield_fixture
+def start_websocket_server(open_port):
+    event_loop = asyncio.new_event_loop()
+
+    def run_server():
+        async def empty_server(websocket, path):
+            data = await websocket.recv()
+            await asyncio.sleep(0.02)
+            await websocket.send(data)
+        server = websockets.serve(empty_server, '127.0.0.1', open_port, loop=event_loop)
+        event_loop.run_until_complete(server)
+        event_loop.run_forever()
+
+    thd = Thread(target=run_server)
+    thd.start()
+    try:
+        yield
+    finally:
+        event_loop.call_soon_threadsafe(event_loop.stop)
+
+
+@pytest.fixture()
+def w3(open_port, start_websocket_server):
+    # need new event loop as the one used by server is already running
+    event_loop = asyncio.new_event_loop()
+    endpoint_uri = 'ws://127.0.0.1:{}'.format(open_port)
+    event_loop.run_until_complete(wait_for_ws(endpoint_uri, event_loop))
+    provider = WebsocketProvider(endpoint_uri, websocket_timeout=0.01)
+    return Web3(provider)
+
+
+def test_websocket_provider_timeout(w3):
+    with pytest.raises(TimeoutError):
+        w3.eth.accounts
+
+
 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()))
diff --git a/tests/generate_go_ethereum_fixture.py b/tests/generate_go_ethereum_fixture.py
@@ -22,6 +22,9 @@
     to_wei,
 )
 
+from tests.utils import (
+    get_open_port,
+)
 from web3 import Web3
 from web3.utils.module_testing.emitter_contract import (
     EMITTER_ABI,
@@ -100,14 +103,6 @@ def tempdir():
         shutil.rmtree(dir_path)
 
 
-def get_open_port():
-    sock = socket.socket()
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return str(port)
-
-
 def get_geth_binary():
     from geth.install import (
         get_executable_path,
diff --git a/tests/integration/generate_fixtures/common.py b/tests/integration/generate_fixtures/common.py
@@ -76,14 +76,6 @@ def tempdir():
         shutil.rmtree(dir_path)
 
 
-def get_open_port():
-    sock = socket.socket()
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return str(port)
-
-
 def get_geth_binary():
     from geth.install import (
         get_executable_path,
diff --git a/tests/integration/generate_fixtures/go_ethereum.py b/tests/integration/generate_fixtures/go_ethereum.py
@@ -11,6 +11,9 @@
 )
 
 import common
+from tests.utils import (
+    get_open_port,
+)
 from web3 import Web3
 from web3.utils.module_testing.emitter_contract import (
     EMITTER_ABI,
@@ -42,7 +45,7 @@ def generate_go_ethereum_fixture(destination_dir):
         geth_ipc_path_dir = stack.enter_context(common.tempdir())
         geth_ipc_path = os.path.join(geth_ipc_path_dir, 'geth.ipc')
 
-        geth_port = common.get_open_port()
+        geth_port = get_open_port()
         geth_binary = common.get_geth_binary()
 
         geth_proc = stack.enter_context(common.get_geth_process(  # noqa: F841
diff --git a/tests/integration/generate_fixtures/parity.py b/tests/integration/generate_fixtures/parity.py
@@ -12,6 +12,9 @@
 
 import common
 import go_ethereum
+from tests.utils import (
+    get_open_port,
+)
 from web3 import Web3
 from web3.utils.toolz import (
     merge,
@@ -176,7 +179,7 @@ def generate_parity_fixture(destination_dir):
 
         geth_datadir = stack.enter_context(common.tempdir())
 
-        geth_port = common.get_open_port()
+        geth_port = get_open_port()
 
         geth_ipc_path_dir = stack.enter_context(common.tempdir())
         geth_ipc_path = os.path.join(geth_ipc_path_dir, 'geth.ipc')
@@ -221,7 +224,7 @@ def generate_parity_fixture(destination_dir):
         parity_ipc_path_dir = stack.enter_context(common.tempdir())
         parity_ipc_path = os.path.join(parity_ipc_path_dir, 'jsonrpc.ipc')
 
-        parity_port = common.get_open_port()
+        parity_port = get_open_port()
         parity_binary = get_parity_binary()
 
         parity_proc = stack.enter_context(get_parity_process(  # noqa: F841
diff --git a/tests/integration/go_ethereum/common.py b/tests/integration/go_ethereum/common.py
@@ -1,5 +1,4 @@
 import pytest
-import socket
 
 from web3.utils.module_testing import (
     EthModuleTest,
@@ -10,14 +9,6 @@
 )
 
 
-def get_open_port():
-    sock = socket.socket()
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return str(port)
-
-
 class GoEthereumTest(Web3ModuleTest):
     def _check_web3_clientVersion(self, client_version):
         assert client_version.startswith('Geth/')
diff --git a/tests/integration/go_ethereum/test_goethereum_http.py b/tests/integration/go_ethereum/test_goethereum_http.py
@@ -1,5 +1,8 @@
 import pytest
 
+from tests.utils import (
+    get_open_port,
+)
 from web3 import Web3
 
 from .common import (
@@ -8,7 +11,6 @@
     GoEthereumPersonalModuleTest,
     GoEthereumTest,
     GoEthereumVersionModuleTest,
-    get_open_port,
 )
 from .utils import (
     wait_for_http,
diff --git a/tests/integration/go_ethereum/test_goethereum_ipc.py b/tests/integration/go_ethereum/test_goethereum_ipc.py
@@ -2,6 +2,9 @@
 import pytest
 import tempfile
 
+from tests.utils import (
+    get_open_port,
+)
 from web3 import Web3
 
 from .common import (
@@ -12,7 +15,6 @@
     GoEthereumVersionModuleTest,
 )
 from .utils import (
-    get_open_port,
     wait_for_socket,
 )
 
diff --git a/tests/integration/go_ethereum/test_goethereum_ws.py b/tests/integration/go_ethereum/test_goethereum_ws.py
@@ -3,7 +3,8 @@
 from tests.integration.common import (
     MiscWebsocketTest,
 )
-from tests.integration.utils import (
+from tests.utils import (
+    get_open_port,
     wait_for_ws,
 )
 from web3 import Web3
@@ -14,7 +15,6 @@
     GoEthereumPersonalModuleTest,
     GoEthereumTest,
     GoEthereumVersionModuleTest,
-    get_open_port,
 )
 
 
diff --git a/tests/integration/go_ethereum/utils.py b/tests/integration/go_ethereum/utils.py
@@ -50,11 +50,3 @@ def kill_proc_gracefully(proc):
     if proc.poll() is None:
         proc.kill()
         wait_for_popen(proc, 2)
-
-
-def get_open_port():
-    sock = socket.socket()
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return str(port)
diff --git a/tests/integration/parity/common.py b/tests/integration/parity/common.py
@@ -1,5 +1,4 @@
 import pytest
-import socket
 
 from flaky import (
     flaky,
@@ -16,14 +15,6 @@
 MAX_FLAKY_RUNS = 3
 
 
-def get_open_port():
-    sock = socket.socket()
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return str(port)
-
-
 class ParityWeb3ModuleTest(Web3ModuleTest):
     def _check_web3_clientVersion(self, client_version):
         assert client_version.startswith('Parity/')
diff --git a/tests/integration/parity/test_parity_http.py b/tests/integration/parity/test_parity_http.py
@@ -4,6 +4,9 @@
 from tests.integration.parity.utils import (
     wait_for_http,
 )
+from tests.utils import (
+    get_open_port,
+)
 from web3 import Web3
 from web3.utils.module_testing import (
     NetModuleTest,
@@ -15,7 +18,6 @@
     ParityPersonalModuleTest,
     ParityTraceModuleTest,
     ParityWeb3ModuleTest,
-    get_open_port,
 )
 
 
diff --git a/tests/integration/parity/test_parity_ws.py b/tests/integration/parity/test_parity_ws.py
@@ -4,7 +4,8 @@
 from tests.integration.common import (
     MiscWebsocketTest,
 )
-from tests.integration.utils import (
+from tests.utils import (
+    get_open_port,
     wait_for_ws,
 )
 from web3 import Web3
@@ -18,7 +19,6 @@
     ParityPersonalModuleTest,
     ParityTraceModuleTest,
     ParityWeb3ModuleTest,
-    get_open_port,
 )
 
 
diff --git a/tests/integration/parity/utils.py b/tests/integration/parity/utils.py
@@ -80,11 +80,3 @@ def kill_proc_gracefully(proc):
     if proc.poll() is None:
         proc.kill()
         wait_for_popen(proc, 2)
-
-
-def get_open_port():
-    sock = socket.socket()
-    sock.bind(('127.0.0.1', 0))
-    port = sock.getsockname()[1]
-    sock.close()
-    return str(port)
diff --git a/tests/utils.py b/tests/utils.py
@@ -1,9 +1,18 @@
 import asyncio
+import socket
 import time
 
 import websockets
 
 
+def get_open_port():
+    sock = socket.socket()
+    sock.bind(('127.0.0.1', 0))
+    port = sock.getsockname()[1]
+    sock.close()
+    return str(port)
+
+
 async def wait_for_ws(endpoint_uri, event_loop, timeout=60):
     start = time.time()
     while time.time() < start + timeout:
diff --git a/web3/providers/websocket.py b/web3/providers/websocket.py
@@ -16,6 +16,7 @@
 )
 
 RESTRICTED_WEBSOCKET_KWARGS = {'uri', 'loop'}
+DEFAULT_WEBSOCKET_TIMEOUT = 10
 
 
 def _start_event_loop(loop):
@@ -63,8 +64,14 @@ class WebsocketProvider(JSONBaseProvider):
     logger = logging.getLogger("web3.providers.WebsocketProvider")
     _loop = None
 
-    def __init__(self, endpoint_uri=None, websocket_kwargs=None):
+    def __init__(
+            self,
+            endpoint_uri=None,
+            websocket_kwargs=None,
+            websocket_timeout=DEFAULT_WEBSOCKET_TIMEOUT
+    ):
         self.endpoint_uri = endpoint_uri
+        self.websocket_timeout = websocket_timeout
         if self.endpoint_uri is None:
             self.endpoint_uri = get_default_endpoint()
         if WebsocketProvider._loop is None:
@@ -90,8 +97,16 @@ def __str__(self):
 
     async def coro_make_request(self, request_data):
         async with self.conn as conn:
-            await conn.send(request_data)
-            return json.loads(await conn.recv())
+            await asyncio.wait_for(
+                conn.send(request_data),
+                timeout=self.websocket_timeout
+            )
+            return json.loads(
+                await asyncio.wait_for(
+                    conn.recv(),
+                    timeout=self.websocket_timeout
+                )
+            )
 
     def make_request(self, method, params):
         self.logger.debug("Making request WebSocket. URI: %s, "

Original file line number	Diff line number	Diff line change
`@@ -2,6 +2,9 @@`
`2`	`2`	`import pytest`
`3`	`3`	`import tempfile`
`4`	`4`
	`5`	`+from tests.utils import (`
	`6`	`+ get_open_port,`
	`7`	`+)`
`5`	`8`	`from web3 import Web3`
`6`	`9`
`7`	`10`	`from .common import (`
`@@ -12,7 +15,6 @@`
`12`	`15`	`GoEthereumVersionModuleTest,`
`13`	`16`	`)`
`14`	`17`	`from .utils import (`
`15`		`- get_open_port,`
`16`	`18`	`wait_for_socket,`
`17`	`19`	`)`
`18`	`20`
Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,8 @@`
`3`	`3`	`from tests.integration.common import (`
`4`	`4`	`MiscWebsocketTest,`
`5`	`5`	`)`
`6`		`-from tests.integration.utils import (`
	`6`	`+from tests.utils import (`
	`7`	`+ get_open_port,`
`7`	`8`	`wait_for_ws,`
`8`	`9`	`)`
`9`	`10`	`from web3 import Web3`
`@@ -14,7 +15,6 @@`
`14`	`15`	`GoEthereumPersonalModuleTest,`
`15`	`16`	`GoEthereumTest,`
`16`	`17`	`GoEthereumVersionModuleTest,`
`17`		`- get_open_port,`
`18`	`18`	`)`
`19`	`19`
`20`	`20`