From 48d1503f6537f9b05d12baa3792db6204f98cbe5 Mon Sep 17 00:00:00 2001
From: Alec Rosenbaum <alec.rosenbaum@close.com>
Date: Tue, 25 Jul 2023 13:44:53 -0400
Subject: [PATCH] Update client to respect the newer `RateLimit` header (#115)

* update test to target new header (currently failing)

* fix test

* bump version
---
 closeio_api/__init__.py | 15 ++++++++++++++-
 tests/test_api.py       | 17 +++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

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.