diff --git a/massive/rest/__init__.py b/massive/rest/__init__.py index 237ab84c..5fe653a3 100644 --- a/massive/rest/__init__.py +++ b/massive/rest/__init__.py @@ -66,6 +66,7 @@ def __init__( ca_certs: Optional[str] = None, proxy: Optional[str] = None, connection_pool_kw: Optional[Dict[str, Any]] = None, + raise_on_status: bool = True, ): """ :param maxsize: Connections kept open per host, for reuse across threads. @@ -75,6 +76,9 @@ def __init__( Environment proxy variables (HTTPS_PROXY etc.) are not used. :param connection_pool_kw: Other urllib3 PoolManager options. They override the client's defaults. + :param raise_on_status: When a retried status (413, 429, 499, 500, 502, 503, + 504) is still returned after the last retry: if True, urllib3 raises + MaxRetryError; if False, the client raises BadResponse with that status. """ # resolve once, so the main and vx clients share it and warn only once ca_certs = ca_certs or ca_bundle_path(logger) @@ -93,6 +97,7 @@ def __init__( ca_certs=ca_certs, proxy=proxy, connection_pool_kw=connection_pool_kw, + raise_on_status=raise_on_status, ) self.vx = VXClient( api_key=api_key, @@ -109,4 +114,5 @@ def __init__( ca_certs=ca_certs, proxy=proxy, connection_pool_kw=connection_pool_kw, + raise_on_status=raise_on_status, ) diff --git a/massive/rest/base.py b/massive/rest/base.py index 7b728685..859ff824 100644 --- a/massive/rest/base.py +++ b/massive/rest/base.py @@ -38,6 +38,7 @@ def __init__( ca_certs: Optional[str] = None, proxy: Optional[str] = None, connection_pool_kw: Optional[Dict[str, Any]] = None, + raise_on_status: bool = True, ): if api_key is None: raise AuthError( @@ -70,6 +71,8 @@ def __init__( 504, ], # default 413, 429, 503 backoff_factor=0.1, # [0.0s, 0.2s, 0.4s, 0.8s, 1.6s, ...] + # if False, return the last response so that _get raises BadResponse + raise_on_status=raise_on_status, ) # https://urllib3.readthedocs.io/en/stable/reference/urllib3.util.html#urllib3.util.Timeout diff --git a/test_rest/test_bad_response.py b/test_rest/test_bad_response.py index 23c86e7b..9bbe79a3 100644 --- a/test_rest/test_bad_response.py +++ b/test_rest/test_bad_response.py @@ -1,3 +1,12 @@ +from http.server import BaseHTTPRequestHandler, HTTPServer +from unittest import mock +import threading +import unittest + +from urllib3.exceptions import MaxRetryError +from urllib3.util.retry import Retry + +from massive import RESTClient from massive.exceptions import BadResponse from base import BaseTest @@ -36,3 +45,67 @@ def test_status_does_not_change_message(self): self.assertEqual(exc.args, ("boom",)) self.assertEqual(str(exc), "boom") self.assertEqual(exc.status, 500) + + +class StatusServer(HTTPServer): + """Answers every GET with a 429 and counts the requests. + + pook replaces urllib3's urlopen, which is where retries happen, so the + raise_on_status tests need a real server. + """ + + def __init__(self): + self.requests = 0 + super().__init__(("127.0.0.1", 0), StatusHandler) + + +class StatusHandler(BaseHTTPRequestHandler): + server: StatusServer + + def do_GET(self): + self.server.requests += 1 + self.send_response(429) + self.send_header("Content-Length", "4") + self.end_headers() + self.wfile.write(b"slow") + + def log_message(self, format, *args): + pass + + +class RaiseOnStatusTest(unittest.TestCase): + def setUp(self): + # if an earlier class left pook on + if pook.isactive(): + pook.disable() + self.addCleanup(pook.activate) + self.server = StatusServer() + threading.Thread( + target=self.server.serve_forever, + kwargs={"poll_interval": 0.01}, + daemon=True, + ).start() + self.addCleanup(self.server.server_close) + self.addCleanup(self.server.shutdown) + + def get_error(self, error, **kwargs): + base = f"http://127.0.0.1:{self.server.server_port}" + c = RESTClient("key", base=base, retries=2, **kwargs) + with mock.patch.object(Retry, "sleep"): # skip the backoff + with self.assertRaises(error) as ctx: + c.get_market_status() + self.assertEqual(self.server.requests, 3) # first attempt + 2 retries + return ctx.exception + + def test_default_raises_max_retry_error(self): + self.get_error(MaxRetryError) + + def test_false_raises_bad_response(self): + exc = self.get_error(BadResponse, raise_on_status=False) + self.assertEqual(exc.status, 429) + self.assertEqual(str(exc), "slow") + + def test_reaches_vx_client(self): + c = RESTClient("key", raise_on_status=False) + for pool in (c.client, c.vx.client): + self.assertFalse(pool.connection_pool_kw["retries"].raise_on_status)