Added automatic retries for specific HTTP response codes. (#6)

- Implemented automatic retries when receiving a 500, 502, or 504 response.
- Improved exception handling logic to handle edge cases

Author: alannix-lw

laceworksdk/exceptions.py

@@ -48,10 +48,11 @@ def __init__(self, response):
             except ValueError:
                 logger.warning("Error parsing JSON response body")
 
-        if 'data' in self.details.keys():
-            self.message = self.details["data"].get("message")
-        elif 'message' in self.details.keys():
-            self.message = self.details["message"]
+        if self.details:
+            if 'data' in self.details.keys():
+                self.message = self.details["data"].get("message")
+            elif 'message' in self.details.keys():
+                self.message = self.details["message"]
         else:
             self.message = None
         """The error message from the parsed API response."""

laceworksdk/http_session.py

@@ -8,6 +8,8 @@
 import requests
 
 from datetime import datetime, timezone
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
 
 from laceworksdk import version
 from laceworksdk.config import (
@@ -43,7 +45,7 @@ def __init__(self, account, subaccount, api_key, api_secret):
         super(HttpSession, self).__init__()
 
         # Create a requests session
-        self._session = requests.Session()
+        self._session = self._retry_session()
 
         # Set the base parameters
         self._api_key = api_key
@@ -54,6 +56,36 @@ def __init__(self, account, subaccount, api_key, api_secret):
         # Get an access token
         self._check_access_token()
 
+    def _retry_session(self,
+                       retries=3,
+                       backoff_factor=0.3,
+                       status_forcelist=(500, 502, 504)):
+        """
+        A method to set up automatic retries on HTTP requests that fail.
+        """
+
+        # Create a new requests session
+        session = requests.Session()
+
+        # Establish the retry criteria
+        retry = Retry(
+            total=retries,
+            read=retries,
+            connect=retries,
+            status=retries,
+            backoff_factor=backoff_factor,
+            status_forcelist=status_forcelist,
+        )
+
+        # Build the adapter with the retry criteria
+        adapter = HTTPAdapter(max_retries=retry)
+
+        # Bind the adapter to HTTP/HTTPS calls
+        session.mount('http://', adapter)
+        session.mount('https://', adapter)
+
+        return session
+
     def _check_access_token(self):
         """
         A method to check the validity of the access token.