Allow timeout for chunked responses (#3428)

* Allow timeout for chunked responses

- Related to #3418

* Only read iter_content if stream is True

* Add newsfragment

---------

Co-authored-by: kclowes <[email protected]>

Author: Uxio0

newsfragments/3428.feature.rst

@@ -0,0 +1 @@
+Add explicit stream kwarg to HTTPProvider so that timeout can be more finely tuned.

tests/core/utilities/test_http_session_manager.py

@@ -29,6 +29,9 @@
 from web3._utils.http_session_manager import (
     HTTPSessionManager,
 )
+from web3.exceptions import (
+    TimeExhausted,
+)
 from web3.utils.caching import (
     SimpleCache,
 )
@@ -44,6 +47,15 @@ def __init__(self, text="", status_code=200):
         self.reason = None
         self.content = "content"
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+    def iter_content(self):
+        return [b"iter content"]
+
     @staticmethod
     def json():
         return json.dumps({"data": "content"})
@@ -110,7 +122,59 @@ def test_session_manager_make_post_request_no_args(mocker, http_session_manager)
     assert len(http_session_manager.session_cache) == 1
     cache_key = generate_cache_key(f"{threading.get_ident()}:{TEST_URI}")
     session = http_session_manager.session_cache.get_cache_entry(cache_key)
-    session.post.assert_called_once_with(TEST_URI, data=b"request", timeout=30)
+    session.post.assert_called_once_with(
+        TEST_URI, data=b"request", timeout=30, stream=False
+    )
+
+    # Ensure the adapter was created with default values
+    check_adapters_mounted(session)
+    adapter = session.get_adapter(TEST_URI)
+    assert isinstance(adapter, HTTPAdapter)
+    assert adapter._pool_connections == DEFAULT_POOLSIZE
+    assert adapter._pool_maxsize == DEFAULT_POOLSIZE
+
+
+def test_session_manager_make_post_request_streaming(mocker, http_session_manager):
+    mocker.patch("requests.Session.post", return_value=MockedResponse())
+
+    # Submit a first request to create a session
+    assert len(http_session_manager.session_cache) == 0
+    response = http_session_manager.make_post_request(
+        TEST_URI, data=b"request", stream=True
+    )
+    assert response == b"iter content"
+    assert len(http_session_manager.session_cache) == 1
+    cache_key = generate_cache_key(f"{threading.get_ident()}:{TEST_URI}")
+    session = http_session_manager.session_cache.get_cache_entry(cache_key)
+    session.post.assert_called_once_with(
+        TEST_URI, data=b"request", timeout=30, stream=True
+    )
+
+    # Ensure the adapter was created with passed in values
+    check_adapters_mounted(session)
+    adapter = session.get_adapter(TEST_URI)
+    assert isinstance(adapter, HTTPAdapter)
+    assert adapter._pool_connections == DEFAULT_POOLSIZE
+    assert adapter._pool_maxsize == DEFAULT_POOLSIZE
+
+
+def test_session_manager_make_post_request_times_out_while_streaming(
+    mocker, http_session_manager
+):
+    mocker.patch("requests.Session.post", return_value=MockedResponse())
+
+    # Submit a first request to create a session
+    assert len(http_session_manager.session_cache) == 0
+    with pytest.raises(TimeExhausted):
+        http_session_manager.make_post_request(
+            TEST_URI, data=b"request", stream=True, timeout=0.000001
+        )
+    assert len(http_session_manager.session_cache) == 1
+    cache_key = generate_cache_key(f"{threading.get_ident()}:{TEST_URI}")
+    session = http_session_manager.session_cache.get_cache_entry(cache_key)
+    session.post.assert_called_once_with(
+        TEST_URI, data=b"request", timeout=0.000001, stream=True
+    )
 
     # Ensure the adapter was created with default values
     check_adapters_mounted(session)
@@ -140,7 +204,9 @@ def test_session_manager_precached_session(mocker, http_session_manager):
 
     # Ensure the timeout was passed to the request
     session = http_session_manager.cache_and_return_session(TEST_URI)
-    session.post.assert_called_once_with(TEST_URI, data=b"request", timeout=60)
+    session.post.assert_called_once_with(
+        TEST_URI, data=b"request", timeout=60, stream=False
+    )
 
     # Ensure the adapter parameters match those we specified
     check_adapters_mounted(session)

web3/_utils/http_session_manager.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import threading
+import time
 from typing import (
     Any,
     Dict,
@@ -32,6 +33,9 @@
 from web3._utils.http import (
     DEFAULT_HTTP_TIMEOUT,
 )
+from web3.exceptions import (
+    TimeExhausted,
+)
 from web3.utils.caching import (
     SimpleCache,
 )
@@ -115,21 +119,40 @@ def json_make_get_request(
     def get_response_from_post_request(
         self, endpoint_uri: URI, *args: Any, **kwargs: Any
     ) -> requests.Response:
-        kwargs.setdefault("timeout", DEFAULT_HTTP_TIMEOUT)
         session = self.cache_and_return_session(
             endpoint_uri, request_timeout=kwargs["timeout"]
         )
-        response = session.post(endpoint_uri, *args, **kwargs)
-        return response
+        return session.post(endpoint_uri, *args, **kwargs)
 
     def make_post_request(
         self, endpoint_uri: URI, data: Union[bytes, Dict[str, Any]], **kwargs: Any
     ) -> bytes:
-        response = self.get_response_from_post_request(
+        kwargs.setdefault("timeout", DEFAULT_HTTP_TIMEOUT)
+        kwargs.setdefault("stream", False)
+
+        start = time.time()
+        timeout = kwargs["timeout"]
+
+        with self.get_response_from_post_request(
             endpoint_uri, data=data, **kwargs
-        )
-        response.raise_for_status()
-        return response.content
+        ) as response:
+            response.raise_for_status()
+            if kwargs.get("stream"):
+                return self._handle_streaming_response(response, start, timeout)
+            else:
+                return response.content
+
+    def _handle_streaming_response(
+        self, response: requests.Response, start: float, timeout: float
+    ) -> bytes:
+        response_body = b""
+        for data in response.iter_content():
+            response_body += data
+            # Manually manage timeout so streaming responses time out
+            # rather than resetting the timeout each time a response comes back
+            if (time.time() - start) > timeout:
+                raise TimeExhausted
+        return response_body
 
     def _close_evicted_sessions(self, evicted_sessions: List[requests.Session]) -> None:
         for evicted_session in evicted_sessions: