From 0a79ad79f61484c56a4ed4af17338d42b8731d07 Mon Sep 17 00:00:00 2001 From: mhorky Date: Mon, 23 Sep 2024 10:35:58 +0200 Subject: [PATCH] fix: Handle Retry-After headers better for 429 responses We have to ensure we normalize the headers before we search for the Retry-After header. --- src/rhsm/connection.py | 8 ++++++-- test/rhsm/unit/test_connection.py | 12 ++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/src/rhsm/connection.py b/src/rhsm/connection.py index 0dc6afcf3..565583acf 100644 --- a/src/rhsm/connection.py +++ b/src/rhsm/connection.py @@ -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 diff --git a/test/rhsm/unit/test_connection.py b/test/rhsm/unit/test_connection.py index 90c750bb2..e57da87a9 100644 --- a/test/rhsm/unit/test_connection.py +++ b/test/rhsm/unit/test_connection.py @@ -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", "")