Adding retry options to auth class

This commit is contained in:
Kevin Fronczak
2020-07-20 01:25:09 +00:00
parent f88dbdaeba
commit a9d3d5568b
2 changed files with 23 additions and 4 deletions
+13 -4
View File
@@ -14,7 +14,7 @@ _LOGGER = logging.getLogger(__name__)
class Auth:
"""Class to handle login communication."""
def __init__(self, login_data=None, no_prompt=False):
def __init__(self, login_data=None, no_prompt=False, retry_opts=None):
"""
Initialize auth handler.
@@ -24,6 +24,10 @@ class Auth:
- password
:param no_prompt: Should any user input prompts
be supressed? True/FALSE
:param retry_opts: Dictionary containing retry options:
- backoff: backoff factor for http request (backoff*(2^(total_retries)-1)
- retries: total retries to attempt
- retry_list: list of status codes to force retry
"""
if login_data is None:
login_data = {}
@@ -36,7 +40,7 @@ class Auth:
self.login_response = None
self.is_errored = False
self.no_prompt = no_prompt
self.session = self.create_session()
self.session = self.create_session(opts=retry_opts)
@property
def login_attributes(self):
@@ -55,15 +59,20 @@ class Auth:
return None
return {"TOKEN_AUTH": self.token}
def create_session(self):
def create_session(self, opts=None):
"""Create a session for blink communication."""
if opts is None:
opts = {}
backoff = opts.get("backoff", 1)
retries = opts.get("retries", 3)
retry_list = opts.get("retry_list", [429, 500, 502, 503, 504])
sess = Session()
assert_status_hook = [
lambda response, *args, **kwargs: response.raise_for_status()
]
sess.hooks["response"] = assert_status_hook
retry = Retry(
total=3, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]
total=retries, backoff_factor=backoff, status_forcelist=retry_list
)
adapter = HTTPAdapter(max_retries=retry)
sess.mount("https://", adapter)
+10
View File
@@ -238,6 +238,16 @@ class TestAuth(unittest.TestCase):
mock_validate.side_effect = [UnauthorizedError, TokenRefreshFailed]
self.assertEqual(self.auth.query(url="http://example.com"), None)
def test_default_session(self):
"""Test default session creation."""
sess = self.auth.create_session()
adapter = sess.adapters["https://"]
self.assertEqual(adapter.max_retries.total, 3)
self.assertEqual(adapter.max_retries.backoff_factor, 1)
self.assertEqual(
adapter.max_retries.status_forcelist, [429, 500, 502, 503, 504]
)
class MockSession:
"""Object to mock a session."""