feat: add raise_on_status to RESTClient

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,
         )

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

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)