diff --git a/descope/auth.py b/descope/auth.py index bf63c66e..92d73d43 100644 --- a/descope/auth.py +++ b/descope/auth.py @@ -105,18 +105,34 @@ def __init__( self.public_keys = {kid: (pub_key, alg)} def _raise_rate_limit_exception(self, response): - resp = response.json() - raise RateLimitException( - resp.get("errorCode", HTTPStatus.TOO_MANY_REQUESTS), - ERROR_TYPE_API_RATE_LIMIT, - resp.get("errorDescription", ""), - resp.get("errorMessage", ""), - rate_limit_parameters={ - API_RATE_LIMIT_RETRY_AFTER_HEADER: int( - response.headers.get(API_RATE_LIMIT_RETRY_AFTER_HEADER, 0) - ) - }, - ) + try: + resp = response.json() + raise RateLimitException( + resp.get("errorCode", HTTPStatus.TOO_MANY_REQUESTS), + ERROR_TYPE_API_RATE_LIMIT, + resp.get("errorDescription", ""), + resp.get("errorMessage", ""), + rate_limit_parameters={ + API_RATE_LIMIT_RETRY_AFTER_HEADER: self._parse_retry_after( + response.headers + ) + }, + ) + except RateLimitException: + raise + except Exception as e: + raise RateLimitException( + status_code=HTTPStatus.TOO_MANY_REQUESTS, + error_type=ERROR_TYPE_API_RATE_LIMIT, + error_message=ERROR_TYPE_API_RATE_LIMIT, + error_description=ERROR_TYPE_API_RATE_LIMIT, + ) + + def _parse_retry_after(self, headers): + try: + return int(headers.get(API_RATE_LIMIT_RETRY_AFTER_HEADER, 0)) + except (ValueError, TypeError): + return 0 def do_get( self, diff --git a/tests/management/test_user.py b/tests/management/test_user.py index 6a009943..86636d82 100644 --- a/tests/management/test_user.py +++ b/tests/management/test_user.py @@ -965,7 +965,7 @@ def test_search_all(self): verify=True, timeout=DEFAULT_TIMEOUT_SECONDS, ) - + # Test success flow with time parameters with patch("requests.post") as mock_post: network_resp = mock.Mock() diff --git a/tests/test_auth.py b/tests/test_auth.py index 40cb1e3b..9de52d72 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -1,3 +1,4 @@ +from http import HTTPStatus import json import os import unittest @@ -678,6 +679,86 @@ def test_api_rate_limit_exception(self): {API_RATE_LIMIT_RETRY_AFTER_HEADER: 10}, ) + def test_api_rate_limit_invalid_header(self): + auth = Auth(self.dummy_project_id, self.public_key_dict) + + # Test do_post empty body + with patch("requests.post") as mock_request: + mock_request.return_value.ok = False + mock_request.return_value.status_code = 429 + mock_request.return_value.json.return_value = { + "errorCode": "E130429", + "errorDescription": "https://docs.descope.com/rate-limit", + "errorMessage": "API rate limit exceeded.", + } + mock_request.return_value.headers = { + API_RATE_LIMIT_RETRY_AFTER_HEADER: "hello" + } + with self.assertRaises(RateLimitException) as cm: + auth.do_post("http://test.com", {}, None, None) + the_exception = cm.exception + self.assertEqual(the_exception.status_code, "E130429") + self.assertEqual(the_exception.error_type, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual( + the_exception.error_description, "https://docs.descope.com/rate-limit" + ) + self.assertEqual(the_exception.error_message, "API rate limit exceeded.") + self.assertEqual( + the_exception.rate_limit_parameters, + {API_RATE_LIMIT_RETRY_AFTER_HEADER: 0}, + ) + + def test_api_rate_limit_invalid_response_body(self): + auth = Auth(self.dummy_project_id, self.public_key_dict) + + # Test do_post empty body + with patch("requests.post") as mock_request: + mock_request.return_value.ok = False + mock_request.return_value.status_code = 429 + mock_request.return_value.json.return_value = "aaa" + with self.assertRaises(RateLimitException) as cm: + auth.do_post("http://test.com", {}, None, None) + the_exception = cm.exception + self.assertEqual(the_exception.status_code, HTTPStatus.TOO_MANY_REQUESTS) + self.assertEqual(the_exception.error_type, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.error_description, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.error_message, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.rate_limit_parameters, {}) + + def test_api_rate_limit_empty_response_body(self): + auth = Auth(self.dummy_project_id, self.public_key_dict) + + # Test do_post empty body + with patch("requests.post") as mock_request: + mock_request.return_value.ok = False + mock_request.return_value.status_code = 429 + mock_request.return_value.json.return_value = "" + with self.assertRaises(RateLimitException) as cm: + auth.do_post("http://test.com", {}, None, None) + the_exception = cm.exception + self.assertEqual(the_exception.status_code, HTTPStatus.TOO_MANY_REQUESTS) + self.assertEqual(the_exception.error_type, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.error_description, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.error_message, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.rate_limit_parameters, {}) + + def test_api_rate_limit_none_response_body(self): + auth = Auth(self.dummy_project_id, self.public_key_dict) + + # Test do_post empty body + with patch("requests.post") as mock_request: + mock_request.return_value.ok = False + mock_request.return_value.status_code = 429 + mock_request.return_value.json.return_value = None + with self.assertRaises(RateLimitException) as cm: + auth.do_post("http://test.com", {}, None, None) + the_exception = cm.exception + self.assertEqual(the_exception.status_code, HTTPStatus.TOO_MANY_REQUESTS) + self.assertEqual(the_exception.error_type, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.error_description, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.error_message, ERROR_TYPE_API_RATE_LIMIT) + self.assertEqual(the_exception.rate_limit_parameters, {}) + def test_raise_from_response(self): auth = Auth(self.dummy_project_id, self.public_key_dict) with patch("requests.get") as mock_request: