Skip to content

Commit 11ac066

Browse files
committed
Provide option to retry calls to ScriptAuthorizer.refresh
1 parent 2753cdf commit 11ac066

File tree

3 files changed

+69
-11
lines changed

3 files changed

+69
-11
lines changed

CHANGES.rst

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,14 @@ Unreleased
99
**Added**
1010

1111
- 301 redirects result in a ``Redirect`` exception.
12+
- Calls to method ``ScriptAuthorizer.refresh`` are handled by a private method,
13+
``ScriptAuthorizer._refresh_with_retries``. The latter provides the option to retry
14+
requests that result in invalid grants.
15+
16+
**Changed**
17+
18+
- Method ``ScriptAuthorizer.two_factor_callback`` can return either a string or a
19+
tuple.
1220

1321
2.2.0 (2021-06-10)
1422
------------------

prawcore/auth.py

Lines changed: 44 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -357,28 +357,61 @@ class ScriptAuthorizer(Authorizer):
357357
AUTHENTICATOR_CLASS = TrustedAuthenticator
358358

359359
def __init__(
360-
self, authenticator, username, password, two_factor_callback=None
360+
self,
361+
authenticator,
362+
username,
363+
password,
364+
two_factor_callback=None,
361365
):
362366
"""Represent a single personal-use authorization to Reddit's API.
363367
364368
:param authenticator: An instance of :class:`TrustedAuthenticator`.
365369
:param username: The Reddit username of one of the application's developers.
366370
:param password: The password associated with ``username``.
367-
:param two_factor_callback: A function that returns OTPs (One-Time
368-
Passcodes), also known as 2FA auth codes. If this function is
369-
provided, prawcore will call it when authenticating.
371+
:param two_factor_callback: (Optional) A function that returns OTPs (One-Time
372+
Passcodes), also known as 2FA auth codes. If provided, this function should
373+
return either a string of six digits or a 3-tuple of the form
374+
``(<OTP>, <DELAY>, <TRIES>)``, where ``<OTP>`` is a string of six
375+
digits, ``<DELAY>`` is an integer that represents the number of seconds
376+
to sleep between invalid authorization attempts, and ``<TRIES>`` is an
377+
integer that represents the maximum number of authorization attempts to
378+
make before an OAuthException is raised.
370379
371380
"""
372-
super(ScriptAuthorizer, self).__init__(authenticator)
381+
super().__init__(authenticator)
373382
self._username = username
374383
self._password = password
375384
self._two_factor_callback = two_factor_callback
376385

386+
def _refresh_with_retries(self, count=1, delay=0, maxcount=1):
387+
if delay > 0:
388+
time.sleep(delay)
389+
additional_kwargs = {}
390+
otp = self._two_factor_callback and self._two_factor_callback()
391+
if otp:
392+
if isinstance(otp, tuple):
393+
if otp[0]:
394+
additional_kwargs["otp"] = otp[0]
395+
else:
396+
additional_kwargs["otp"] = otp
397+
try:
398+
self._request_token(
399+
grant_type="password",
400+
username=self._username,
401+
password=self._password,
402+
**additional_kwargs,
403+
)
404+
except OAuthException:
405+
if otp and isinstance(otp, tuple) and len(otp) == 3:
406+
_, delay, maxcount = otp
407+
if count >= min(maxcount, 10):
408+
raise
409+
self._refresh_with_retries(
410+
count=count + 1, delay=delay, maxcount=maxcount
411+
)
412+
else:
413+
raise
414+
377415
def refresh(self):
378416
"""Obtain a new personal-use script type access token."""
379-
self._request_token(
380-
grant_type="password",
381-
username=self._username,
382-
password=self._password,
383-
otp=self._two_factor_callback and self._two_factor_callback(),
384-
)
417+
self._refresh_with_retries()

tests/test_authorizer.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import unittest
33

44
from betamax import Betamax
5+
from mock import Mock, patch
56

67
import prawcore
78

@@ -384,3 +385,19 @@ def test_refresh__with_invalid_username_or_password(self):
384385
):
385386
self.assertRaises(prawcore.OAuthException, authorizer.refresh)
386387
self.assertFalse(authorizer.is_valid())
388+
389+
@patch("time.sleep", return_value=None)
390+
@patch("prawcore.Requestor.request")
391+
def test_refresh_with_retries(self, mock_post, _):
392+
response = Mock(
393+
json=lambda: {"error": "invalid grant"}, status_code=200
394+
)
395+
mock_post.side_effect = [response, response]
396+
authorizer = prawcore.ScriptAuthorizer(
397+
self.authentication,
398+
"dummy",
399+
"dummy",
400+
two_factor_callback=lambda: ("123456", 31, 2),
401+
)
402+
with self.assertRaises(prawcore.OAuthException):
403+
authorizer.refresh()

0 commit comments

Comments
 (0)