Skip to content

Commit

Permalink
fix: Handle Retry-After headers better for 429 responses
Browse files Browse the repository at this point in the history
We have to ensure we normalize the headers before we search for the
Retry-After header.
  • Loading branch information
m-horky committed Sep 23, 2024
1 parent dc31ee3 commit 0a79ad7
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 2 deletions.
8 changes: 6 additions & 2 deletions src/rhsm/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -500,10 +500,14 @@ class RateLimitExceededException(RestlibException):
The retry_after attribute may not be included in the response.
"""

def __init__(self, code: int, msg: str = None, headers: str = None) -> None:
def __init__(self, code: int, msg: str = None, headers: dict = None) -> None:
super(RateLimitExceededException, self).__init__(code, msg)
self.headers = headers or {}
self.retry_after = safe_int(self.headers.get("retry-after"))
self.retry_after = None
for header, value in self.headers.items():
if header.lower() == "retry-after":
self.retry_after = safe_int(value)
break
self.msg = msg or "Access rate limit exceeded"
if self.retry_after is not None:
self.msg += ", retry access after: %s seconds." % self.retry_after
Expand Down
12 changes: 12 additions & 0 deletions test/rhsm/unit/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -871,6 +871,18 @@ def test_429_body(self):
else:
self.fail("Should have raised a RateLimitExceededException")

def test_429_weird_case(self):
content = '{"errors": ["TooFast"]}'
headers = {"RETry-aFteR": 20}
try:
self.vr("429", content, headers)
except RateLimitExceededException as e:
self.assertEqual(20, e.retry_after)
self.assertEqual("TooFast, retry access after: 20 seconds.", e.msg)
self.assertEqual("429", e.code)
else:
self.fail("Should have raised a RateLimitExceededException")

def test_500_empty(self):
try:
self.vr("500", "")
Expand Down

0 comments on commit 0a79ad7

Please sign in to comment.