diff --git a/closeio_api/__init__.py b/closeio_api/__init__.py index 331df1c..f0b912f 100644 --- a/closeio_api/__init__.py +++ b/closeio_api/__init__.py @@ -1,5 +1,6 @@ import contextlib import logging +import re import time from random import uniform @@ -12,7 +13,7 @@ # To update the package version, change this variable. This variable is also # read by setup.py when installing the package. -__version__ = '2.0' +__version__ = '2.1' class APIError(Exception): """Raised when sending a request to the API failed.""" @@ -146,8 +147,20 @@ def _get_rate_limit_sleep_time(self, response): """Get rate limit window expiration time from response if the response status code is 429. """ + with contextlib.suppress(KeyError): + rate_limit = response.headers["RateLimit"] + # we don't actually need all these values, but per the RFC: + # "Malformed RateLimit header fields MUST be ignored." + match = re.match( + r"limit=(\d+), remaining=(\d+), reset=(\d+)", rate_limit + ) + if match: + limit, remaining, reset = match.groups() + return float(reset) + with contextlib.suppress(KeyError): return float(response.headers["Retry-After"]) + with contextlib.suppress(KeyError): return float(response.headers["RateLimit-Reset"]) diff --git a/tests/test_api.py b/tests/test_api.py index ed74103..eec88ca 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -107,7 +107,20 @@ def request_callback(request): assert resp['data'][0]['name'] == 'Sample Lead' @responses.activate -def test_retry_on_rate_limit(api_client): +@pytest.mark.parametrize( + "headers", + [ + {"RateLimit-Reset": "1"}, + {"Retry-After": "1"}, + {"RateLimit": "limit=100, remaining=0, reset=1"}, + { + "Retry-After": "1", + "RateLimit-Reset": "1", + "RateLimit": "limit=100, remaining=0, reset=1", + }, + ] +) +def test_retry_on_rate_limit(api_client, headers): with responses.RequestsMock() as rsps: # Rate limit the first request and suggest it can be retried in 1 sec. @@ -117,7 +130,7 @@ def test_retry_on_rate_limit(api_client): body=json.dumps({}), status=429, content_type='application/json', - headers={"Retry-After": "1", "RateLimit-Reset": "1"} + headers=headers, ) # Respond correctly to the second request.