Merge pull request #8 from GoodRx/simpler-retries

Simplify API retries while obeying rate limits

Author: Sam Park

braze/client.py

@@ -1,22 +1,59 @@
-import json
 import time
 
 import requests
 from requests.exceptions import RequestException
 from tenacity import retry
-from tenacity import retry_if_exception_type
 from tenacity import stop_after_attempt
 from tenacity import wait_random_exponential
 
+DEFAULT_API_URL = "https://rest.iad-02.braze.com"
+USER_TRACK_ENDPOINT = "/users/track"
+USER_DELETE_ENDPOINT = "/users/delete"
+MAX_RETRIES = 3
+# Max time to wait between API call retries
+MAX_WAIT_SECONDS = 1.25
+
 
 class BrazeRateLimitError(Exception):
-    pass
+    def __init__(self, reset_epoch_s):
+        """
+        A rate limit error was encountered.
+
+        :param float reset_epoch_s: Unix timestamp for when the API may be called again.
+        """
+        self.reset_epoch_s = reset_epoch_s
+        super(BrazeRateLimitError, self).__init__("BrazeRateLimitError")
 
 
 class BrazeInternalServerError(Exception):
     pass
 
 
+def _wait_random_exp_or_rate_limit():
+    """Creates a tenacity wait callback that accounts for explicit rate limits."""
+    random_exp = wait_random_exponential(multiplier=1, max=MAX_WAIT_SECONDS)
+
+    def check(retry_state):
+        """
+        Waits with either a random exponential backoff or attempts to obey rate limits
+        that Braze returns.
+
+        :param tenacity.RetryCallState retry_state: Info about current retry invocation
+        :raises BrazeRateLimitError: If the rate limit reset time is too long
+        :returns: Time to wait, in seconds.
+        :rtype: float
+        """
+        exc = retry_state.outcome.exception()
+        if isinstance(exc, BrazeRateLimitError):
+            sec_to_reset = exc.reset_epoch_s - float(time.time())
+            if sec_to_reset >= MAX_WAIT_SECONDS:
+                raise exc
+            return max(0.0, sec_to_reset)
+        return random_exp(retry_state=retry_state)
+
+    return check
+
+
 class BrazeClient(object):
     """
     Client for Appboy public API. Support user_track.
@@ -42,15 +79,9 @@ class BrazeClient(object):
         print r['errors']
     """
 
-    DEFAULT_API_URL = "https://rest.iad-02.braze.com"
-    USER_TRACK_ENDPOINT = "/users/track"
-    USER_DELETE_ENDPOINT = "/users/delete"
-    MAX_RETRIES = 3
-    MAX_WAIT_SECONDS = 1.25
-
     def __init__(self, api_key, api_url=None):
         self.api_key = api_key
-        self.api_url = api_url or self.DEFAULT_API_URL
+        self.api_url = api_url or DEFAULT_API_URL
         self.request_url = ""
 
     def user_track(self, attributes, events, purchases):
@@ -61,7 +92,7 @@ def user_track(self, attributes, events, purchases):
         :param purchases: dict or list of user purchases dict (external_id, app_id, product_id, currency, price)
         :return: json dict response, for example: {"message": "success", "errors": [], "client_error": ""}
         """
-        self.request_url = self.api_url + self.USER_TRACK_ENDPOINT
+        self.request_url = self.api_url + USER_TRACK_ENDPOINT
 
         payload = {}
 
@@ -83,7 +114,7 @@ def user_delete(self, external_ids, appboy_ids):
         :param appboy_ids: dict or list of user braze ids
         :return: json dict response, for example: {"message": "success", "errors": [], "client_error": ""}
         """
-        self.request_url = self.api_url + self.USER_DELETE_ENDPOINT
+        self.request_url = self.api_url + USER_DELETE_ENDPOINT
 
         payload = {}
 
@@ -129,35 +160,19 @@ def __create_request(self, payload):
 
     @retry(
         reraise=True,
-        wait=wait_random_exponential(multiplier=1, max=MAX_WAIT_SECONDS),
+        wait=_wait_random_exp_or_rate_limit(),
         stop=stop_after_attempt(MAX_RETRIES),
-        retry=(
-            retry_if_exception_type(BrazeInternalServerError)
-            | retry_if_exception_type(RequestException)
-        ),
     )
-    def _post_request_with_retries(self, payload, retry_attempt=0):
+    def _post_request_with_retries(self, payload):
         """
         :param dict payload:
-        :param int retry_attempt: current retry attempt number
-        :rtype: dict
+        :rtype: requests.Response
         """
-        headers = {"Content-Type": "application/json"}
-        r = requests.post(
-            self.request_url, data=json.dumps(payload), headers=headers, timeout=2
-        )
-        if retry_attempt >= self.MAX_RETRIES:
-            raise BrazeRateLimitError("BrazeRateLimitError")
-
+        r = requests.post(self.request_url, json=payload, timeout=2)
         # https://www.braze.com/docs/developer_guide/rest_api/messaging/#fatal-errors
         if r.status_code == 429:
-            reset_epoch_seconds = float(r.headers.get("X-RateLimit-Reset"))
-            sec_to_reset = reset_epoch_seconds - float(time.time())
-            if sec_to_reset < self.MAX_WAIT_SECONDS:
-                time.sleep(sec_to_reset)
-                return self._post_request_with_retries(payload, retry_attempt + 1)
-            else:
-                raise BrazeRateLimitError("BrazeRateLimitError")
+            reset_epoch_s = float(r.headers.get("X-RateLimit-Reset", 0))
+            raise BrazeRateLimitError(reset_epoch_s)
         elif str(r.status_code).startswith("5"):
             raise BrazeInternalServerError("BrazeInternalServerError")
         return r

setup.py

@@ -2,9 +2,9 @@
 from setuptools import setup
 
 NAME = "braze-client"
-VERSION = "2.0.0"
+VERSION = "2.1.0"
 
-REQUIRES = ["requests >=2.21.0, <3.0.0", "tenacity >=4.8.0, <6.0.0"]
+REQUIRES = ["requests >=2.21.0, <3.0.0", "tenacity >=5.0.0, <6.0.0"]
 
 EXTRAS = {"dev": ["tox"]}
 

tests/braze/test_client.py

@@ -1,10 +1,17 @@
 import time
 
+from braze.client import _wait_random_exp_or_rate_limit
 from braze.client import BrazeClient
+from braze.client import BrazeRateLimitError
+from braze.client import MAX_RETRIES
+from braze.client import MAX_WAIT_SECONDS
 from freezegun import freeze_time
 import pytest
+from pytest import approx
 from requests import RequestException
 from requests_mock import ANY
+from tenacity import Future
+from tenacity import RetryCallState
 
 
 @pytest.fixture
@@ -27,6 +34,42 @@ def purchases():
     return {"external_id": "123", "name": "some_name"}
 
 
+class TestWaitRandomExpOrRateLimit(object):
+    @pytest.fixture
+    def retry_state(self):
+        retry_state = RetryCallState(object(), lambda x: x, (), {})
+        retry_state.outcome = Future(attempt_number=1)
+        return retry_state
+
+    @freeze_time()
+    def test_raises_if_too_long(self, retry_state):
+        callback = _wait_random_exp_or_rate_limit()
+        exc = BrazeRateLimitError(time.time() + MAX_WAIT_SECONDS + 1)
+        retry_state.outcome.set_exception(exc)
+
+        with pytest.raises(BrazeRateLimitError) as e:
+            callback(retry_state)
+
+        assert e.value.reset_epoch_s == exc.reset_epoch_s
+
+    @freeze_time()
+    def test_doesnt_allow_negative_waits(self, retry_state):
+        callback = _wait_random_exp_or_rate_limit()
+        exc = BrazeRateLimitError(time.time() - 1)
+        retry_state.outcome.set_exception(exc)
+
+        assert callback(retry_state) == 0.0
+
+    def test_uses_random_exp_for_other_exceptions(self, retry_state):
+        callback = _wait_random_exp_or_rate_limit()
+        retry_state.outcome.set_exception(Exception())
+
+        for attempt in range(10):
+            retry_state.attempt_number = attempt
+            for _ in range(100):
+                assert 0 <= callback(retry_state) <= 1.5
+
+
 class TestBrazeClient(object):
     def test_init(self, braze_client):
         assert braze_client.api_key == "API_KEY"
@@ -65,7 +108,7 @@ def test_user_track_request_exception(
         assert error_msg in response["errors"]
 
     @pytest.mark.parametrize(
-        "status_code, retry_attempts", [(500, BrazeClient.MAX_RETRIES), (401, 1)]
+        "status_code, retry_attempts", [(500, MAX_RETRIES), (401, 1)]
     )
     def test_retries_for_errors(
         self,
@@ -95,18 +138,18 @@ def test_retries_for_errors(
     @freeze_time()
     @pytest.mark.parametrize(
         "reset_delta_seconds, expected_attempts",
-        [(0.05, 1 + BrazeClient.MAX_RETRIES), (BrazeClient.MAX_WAIT_SECONDS + 1, 1)],
+        [(0.05, MAX_RETRIES), (MAX_WAIT_SECONDS + 1, 1)],
     )
     def test_retries_for_rate_limit_errors(
         self,
         braze_client,
-        mocker,
         requests_mock,
         attributes,
         events,
         purchases,
         reset_delta_seconds,
         expected_attempts,
+        no_sleep,
     ):
         headers = {
             "Content-Type": "application/json",
@@ -116,10 +159,15 @@ def test_retries_for_rate_limit_errors(
         mock_json = {"message": error_msg, "errors": error_msg}
         requests_mock.post(ANY, json=mock_json, status_code=429, headers=headers)
 
-        spy = mocker.spy(braze_client, "_post_request_with_retries")
         response = braze_client.user_track(
             attributes=attributes, events=events, purchases=purchases
         )
-        assert spy.call_count == expected_attempts
+
+        stats = braze_client._post_request_with_retries.retry.statistics
+        assert stats["attempt_number"] == expected_attempts
         assert response["success"] is False
         assert "BrazeRateLimitError" in response["errors"]
+
+        # Ensure the correct wait time is used when rate limited
+        for i in range(expected_attempts - 1):
+            assert approx(no_sleep.call_args_list[i][0], reset_delta_seconds)

tests/conftest.py

@@ -7,3 +7,11 @@
 @pytest.fixture
 def braze_client():
     return BrazeClient(api_key="API_KEY")
+
+
+@pytest.fixture(autouse=True)
+def no_sleep(mocker, braze_client):
+    """Disables actual sleeps, but keeps retry wait logic. Zippy tests!"""
+    return mocker.patch.object(
+        braze_client._post_request_with_retries.retry, "sleep", return_value=None
+    )