File tree Expand file tree Collapse file tree 2 files changed +29
-3
lines changed
Expand file tree Collapse file tree 2 files changed +29
-3
lines changed Original file line number Diff line number Diff line change 11import contextlib
22import logging
3+ import re
34import time
45
56from random import uniform
1213
1314# To update the package version, change this variable. This variable is also
1415# read by setup.py when installing the package.
15- __version__ = '2.0 '
16+ __version__ = '2.1 '
1617
1718class APIError (Exception ):
1819 """Raised when sending a request to the API failed."""
@@ -146,8 +147,20 @@ def _get_rate_limit_sleep_time(self, response):
146147 """Get rate limit window expiration time from response if the response
147148 status code is 429.
148149 """
150+ with contextlib .suppress (KeyError ):
151+ rate_limit = response .headers ["RateLimit" ]
152+ # we don't actually need all these values, but per the RFC:
153+ # "Malformed RateLimit header fields MUST be ignored."
154+ match = re .match (
155+ r"limit=(\d+), remaining=(\d+), reset=(\d+)" , rate_limit
156+ )
157+ if match :
158+ limit , remaining , reset = match .groups ()
159+ return float (reset )
160+
149161 with contextlib .suppress (KeyError ):
150162 return float (response .headers ["Retry-After" ])
163+
151164 with contextlib .suppress (KeyError ):
152165 return float (response .headers ["RateLimit-Reset" ])
153166
Original file line number Diff line number Diff line change @@ -107,7 +107,20 @@ def request_callback(request):
107107 assert resp ['data' ][0 ]['name' ] == 'Sample Lead'
108108
109109@responses .activate
110- def test_retry_on_rate_limit (api_client ):
110+ @pytest .mark .parametrize (
111+ "headers" ,
112+ [
113+ {"RateLimit-Reset" : "1" },
114+ {"Retry-After" : "1" },
115+ {"RateLimit" : "limit=100, remaining=0, reset=1" },
116+ {
117+ "Retry-After" : "1" ,
118+ "RateLimit-Reset" : "1" ,
119+ "RateLimit" : "limit=100, remaining=0, reset=1" ,
120+ },
121+ ]
122+ )
123+ def test_retry_on_rate_limit (api_client , headers ):
111124 with responses .RequestsMock () as rsps :
112125
113126 # 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):
117130 body = json .dumps ({}),
118131 status = 429 ,
119132 content_type = 'application/json' ,
120- headers = { "Retry-After" : "1" , "RateLimit-Reset" : "1" }
133+ headers = headers ,
121134 )
122135
123136 # Respond correctly to the second request.
You can’t perform that action at this time.
0 commit comments