Add support for multiple ratelimit headers (#311)

Author: nsaje

socketshark/config_defaults.py

@@ -52,7 +52,8 @@
     'wait': 3,
     # If we encounter HTTP 429, wait for the amount of seconds specified in
     # the following header (fall back to the 'wait' value if it doesn't exist).
-    'rate_limit_reset_header_name': 'X-Rate-Limit-Reset',
+    'rate_limit_reset_header_name': 'X-Rate-Limit-Reset',  # Deprecated, use rate_limit_reset_header_names instead.
+    'rate_limit_reset_header_names': [],  # List of header names to check for the rate limit reset time.
 }
 
 # Redis options

socketshark/utils.py

@@ -21,9 +21,15 @@ def _get_rate_limit_wait(log, resp, opts):
 
     wait = opts['wait']
 
-    header_name = opts['rate_limit_reset_header_name']
-    if header_name and header_name in resp.headers:
-        header_value = resp.headers[header_name]
+    header_names = opts['rate_limit_reset_header_names'] or [
+        opts['rate_limit_reset_header_name']
+    ]
+    header_value = None
+    for header_name in header_names:
+        if header_name in resp.headers:
+            header_value = resp.headers[header_name]
+            break
+    if header_value:
         try:
             new_wait = float(header_value)
             # Make sure we have a valid value (not negative, NaN, or Inf)
@@ -48,6 +54,10 @@ def _get_rate_limit_wait(log, resp, opts):
                 name=header_name,
                 value=header_value,
             )
+    else:
+        log.warn(
+            'got a 429 but no rate limit reset header found in response',
+        )
     return wait
 
 

tests/test_basic.py

@@ -1855,13 +1855,33 @@ def dummy_ping(*args, **kwargs):
         await shark.shutdown()
 
     @pytest.mark.asyncio
-    async def test_rate_limit(self):
+    @pytest.mark.parametrize(
+        'response_header_name, config_header_names, config_header_name',
+        (
+            (
+                'X-Rate-Limit-Reset',
+                ['X-Rate-Limit-Reset', 'Retry-After'],
+                None,
+            ),
+            ('Retry-After', ['X-Rate-Limit-Reset', 'Retry-After'], None),
+            ('X-Rate-Limit-Reset', [], 'X-Rate-Limit-Reset'),
+        ),
+    )
+    async def test_rate_limit(
+        self, response_header_name, config_header_names, config_header_name
+    ):
         """
         Make sure SocketShark retries 429 responses appropriately.
         """
         http_retry_config = TEST_CONFIG.copy()
         http_retry_config['HTTP']['tries'] = 2
         http_retry_config['HTTP']['wait'] = 1
+        http_retry_config['HTTP'][
+            'rate_limit_reset_header_names'
+        ] = config_header_names
+        http_retry_config['HTTP'][
+            'rate_limit_reset_header_name'
+        ] = config_header_name
         shark = SocketShark(http_retry_config)
         await shark.prepare()
         client = MockClient(shark)
@@ -1876,7 +1896,7 @@ async def test_rate_limit(self):
                 conf['before_subscribe'],
                 status=429,
                 headers={
-                    'X-Rate-Limit-Reset': '0.2',
+                    response_header_name: '0.2',
                 },
             )
             mock.post(