Fix rate limit parsing

Author: dorsha

.vscode/settings.json

@@ -4,5 +4,10 @@
   "flake8.importStrategy": "fromEnvironment",
   "mypy-type-checker.importStrategy": "fromEnvironment",
   "isort.importStrategy": "fromEnvironment",
-  "black-formatter.importStrategy": "fromEnvironment"
+  "black-formatter.importStrategy": "fromEnvironment",
+  "workbench.colorCustomizations": {
+    "activityBar.background": "#121D85",
+    "titleBar.activeBackground": "#1A28BA",
+    "titleBar.activeForeground": "#F9FAFE"
+  }
 }

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,

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()

tests/test_auth.py

@@ -1,3 +1,4 @@
+from http import HTTPStatus
 import json
 import os
 import unittest
@@ -678,6 +679,57 @@ def test_api_rate_limit_exception(self):
                 {API_RATE_LIMIT_RETRY_AFTER_HEADER: 10},
             )
 
+    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: