Provide option to retry calls to ScriptAuthorizer.refresh

Author: MaybeNetwork

CHANGES.rst

@@ -12,6 +12,8 @@ Unreleased
 - ``Requestor`` is now initialzed with a ``timeout`` parameter.
 - ``ScriptAuthorizer``, ``ReadOnlyAuthorizer``, and ``DeviceIDAuthorizer`` have a
   new parameter, ``scopes``, which determines the scope of access requests.
+- ``ScriptAuthorizer.refresh`` can now retry authorization attempts that raise
+  instances of ``OAuthException``.
 
 2.2.0 (2021-06-10)
 ------------------

prawcore/auth.py

@@ -392,28 +392,32 @@ def __init__(
         password,
         two_factor_callback=None,
         scopes=None,
+        retries=0,
     ):
         """Represent a single personal-use authorization to Reddit's API.
 
         :param authenticator: An instance of :class:`TrustedAuthenticator`.
         :param username: The Reddit username of one of the application's developers.
         :param password: The password associated with ``username``.
-        :param two_factor_callback: A function that returns OTPs (One-Time
-            Passcodes), also known as 2FA auth codes. If this function is
-            provided, prawcore will call it when authenticating.
+        :param two_factor_callback: (Optional) A function that returns a two factor
+            authentication code.
         :param scopes: (Optional) A list of OAuth scopes to request authorization for
             (default: None). The scope ``*`` is requested when the default argument is
             used.
+        :param retries: (Optional) The number of times to retry an authorization
+            attempt that raises an ``OAuthException`` (default: 0). The argument should be a
+            nonnegative integer less than or equal to ten. This setting is ignored if
+            two_factor_callback does not return a value that is ``True``.
 
         """
-        super(ScriptAuthorizer, self).__init__(authenticator)
+        super().__init__(authenticator)
         self._password = password
+        self._retries = abs(retries)
         self._scopes = scopes
         self._two_factor_callback = two_factor_callback
         self._username = username
 
-    def refresh(self):
-        """Obtain a new personal-use script type access token."""
+    def _refresh_with_retries(self, count=0):
         additional_kwargs = {}
         if self._scopes:
             additional_kwargs["scope"] = " ".join(self._scopes)
@@ -422,9 +426,21 @@ def refresh(self):
         )
         if two_factor_code:
             additional_kwargs["otp"] = two_factor_code
-        self._request_token(
-            grant_type="password",
-            username=self._username,
-            password=self._password,
-            **additional_kwargs,
-        )
+        try:
+            self._request_token(
+                grant_type="password",
+                username=self._username,
+                password=self._password,
+                **additional_kwargs,
+            )
+        except OAuthException:
+            if two_factor_code:
+                if count >= min(self._retries, 10):
+                    raise
+                self._refresh_with_retries(count + 1)
+            else:
+                raise
+
+    def refresh(self):
+        """Obtain a new personal-use script type access token."""
+        self._refresh_with_retries()

tests/test_authorizer.py

@@ -1,7 +1,9 @@
 """Test for prawcore.auth.Authorizer classes."""
 import unittest
+from traceback import format_exc
 
 from betamax import Betamax
+from mock import Mock, patch
 
 import prawcore
 
@@ -14,8 +16,8 @@
     REFRESH_TOKEN,
     REQUESTOR,
     TEMPORARY_GRANT_CODE,
-    two_factor_callback,
     USERNAME,
+    two_factor_callback,
 )
 
 
@@ -386,7 +388,7 @@ def test_refresh__with_invalid_otp(self):
             "ScriptAuthorizer_refresh__with_invalid_otp"
         ):
             self.assertRaises(prawcore.OAuthException, authorizer.refresh)
-        self.assertFalse(authorizer.is_valid())
+            self.assertFalse(authorizer.is_valid())
 
     def test_refresh__with_invalid_username_or_password(self):
         authorizer = prawcore.ScriptAuthorizer(
@@ -398,6 +400,27 @@ def test_refresh__with_invalid_username_or_password(self):
             self.assertRaises(prawcore.OAuthException, authorizer.refresh)
         self.assertFalse(authorizer.is_valid())
 
+    @patch("time.sleep", return_value=None)
+    @patch("prawcore.Requestor.request")
+    def test_refresh__with_retries(self, mock_post, _):
+        response = Mock(
+            json=lambda: {"error": "invalid grant"}, status_code=200
+        )
+        mock_post.side_effect = [response for x in range(18)]
+        authorizer = prawcore.ScriptAuthorizer(
+            self.authentication,
+            "dummy",
+            "dummy",
+            retries=13,
+            two_factor_callback=lambda: "123456",
+        )
+        try:
+            authorizer.refresh()
+        except prawcore.OAuthException:
+            traceback = format_exc()
+        assert traceback.count("prawcore.exceptions.OAuthException") == 11
+        assert mock_post.call_count == 11
+
     def test_refresh__with_scopes(self):
         scope_list = ["adsedit", "adsread", "creddits", "history"]
         authorizer = prawcore.ScriptAuthorizer(