Skip to content

Commit

Permalink
Add 'on_error' argument to 'retry.retry_target' and 'Retry'. (#3891)
Browse files Browse the repository at this point in the history
Permit application code to reset / fix-up state before retrying.
  • Loading branch information
tseaver authored Aug 29, 2017
1 parent 29654a6 commit 9c4144b
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 6 deletions.
16 changes: 13 additions & 3 deletions core/google/api/core/retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ def exponential_sleep_generator(
delay = delay * multiplier


def retry_target(target, predicate, sleep_generator, deadline):
def retry_target(target, predicate, sleep_generator, deadline, on_error=None):
"""Call a function and retry if it fails.
This is the lowest-level retry helper. Generally, you'll use the
Expand All @@ -150,6 +150,9 @@ def retry_target(target, predicate, sleep_generator, deadline):
sleep_generator (Iterator[float]): An infinite iterator that determines
how long to sleep between retries.
deadline (float): How long to keep retrying the target.
on_error (Callable): A function to call while processing a retryable
exception. Any error raised by this function will *not* be
caught.
Returns:
Any: the return value of the target function.
Expand Down Expand Up @@ -177,6 +180,8 @@ def retry_target(target, predicate, sleep_generator, deadline):
if not predicate(exc):
raise
last_exc = exc
if on_error is not None:
on_error(exc)

now = datetime_helpers.utcnow()
if deadline_datetime is not None and deadline_datetime < now:
Expand Down Expand Up @@ -226,11 +231,14 @@ def __init__(
self._maximum = maximum
self._deadline = deadline

def __call__(self, func):
def __call__(self, func, on_error=None):
"""Wrap a callable with retry behavior.
Args:
func (Callable): The callable to add retry behavior to.
on_error (Callable): A function to call while processing a
retryable exception. Any error raised by this function will
*not* be caught.
Returns:
Callable: A callable that will invoke ``func`` with retry
Expand All @@ -246,7 +254,9 @@ def retry_wrapped_func(*args, **kwargs):
target,
self._predicate,
sleep_generator,
self._deadline)
self._deadline,
on_error=on_error,
)

return retry_wrapped_func

Expand Down
41 changes: 38 additions & 3 deletions core/tests/unit/api_core/test_retry.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,34 @@ def target():
sleep.assert_has_calls([mock.call(0), mock.call(1)])


@mock.patch('time.sleep', autospec=True)
@mock.patch(
'google.api.core.helpers.datetime_helpers.utcnow',
return_value=datetime.datetime.min,
autospec=True)
def test_retry_target_w_on_error(utcnow, sleep):
predicate = retry.if_exception_type(ValueError)
call_count = {'target': 0}
to_raise = ValueError()

def target():
call_count['target'] += 1
if call_count['target'] < 3:
raise to_raise
return 42

on_error = mock.Mock()

result = retry.retry_target(
target, predicate, range(10), None, on_error=on_error)

assert result == 42
assert call_count['target'] == 3

on_error.assert_has_calls([mock.call(to_raise), mock.call(to_raise)])
sleep.assert_has_calls([mock.call(0), mock.call(1)])


@mock.patch('time.sleep', autospec=True)
@mock.patch(
'google.api.core.helpers.datetime_helpers.utcnow',
Expand Down Expand Up @@ -139,7 +167,8 @@ def test_constructor_options(self):
initial=1,
maximum=2,
multiplier=3,
deadline=4)
deadline=4,
)
assert retry_._predicate == mock.sentinel.predicate
assert retry_._initial == 1
assert retry_._maximum == 2
Expand Down Expand Up @@ -204,12 +233,17 @@ def test___call___and_execute_success(self, sleep):
'random.uniform', autospec=True, side_effect=lambda m, n: n/2.0)
@mock.patch('time.sleep', autospec=True)
def test___call___and_execute_retry(self, sleep, uniform):
retry_ = retry.Retry(predicate=retry.if_exception_type(ValueError))

on_error = mock.Mock(spec=['__call__'], side_effect=[None])
retry_ = retry.Retry(
predicate=retry.if_exception_type(ValueError),
)

target = mock.Mock(spec=['__call__'], side_effect=[ValueError(), 42])
# __name__ is needed by functools.partial.
target.__name__ = 'target'

decorated = retry_(target)
decorated = retry_(target, on_error=on_error)
target.assert_not_called()

result = decorated('meep')
Expand All @@ -218,3 +252,4 @@ def test___call___and_execute_retry(self, sleep, uniform):
assert target.call_count == 2
target.assert_has_calls([mock.call('meep'), mock.call('meep')])
sleep.assert_called_once_with(retry_._initial)
assert on_error.call_count == 1

0 comments on commit 9c4144b

Please sign in to comment.