Support custom wait and stop functions.
This commit is contained in:
parent
6d8ef5bd03
commit
8b73aff04e
14
retrying.py
14
retrying.py
@ -108,7 +108,9 @@ class Retrying(object):
|
||||
wait_exponential_multiplier=None, wait_exponential_max=None,
|
||||
retry_on_exception=None,
|
||||
retry_on_result=None,
|
||||
wrap_exception=False):
|
||||
wrap_exception=False,
|
||||
stop_func=None,
|
||||
wait_func=None):
|
||||
|
||||
self._stop_max_attempt_number = 5 if stop_max_attempt_number is None else stop_max_attempt_number
|
||||
self._stop_max_delay = 100 if stop_max_delay is None else stop_max_delay
|
||||
@ -129,7 +131,10 @@ class Retrying(object):
|
||||
if stop_max_delay is not None:
|
||||
stop_funcs.append(self.stop_after_delay)
|
||||
|
||||
if stop is None:
|
||||
if stop_func is not None:
|
||||
self.stop = stop_func
|
||||
|
||||
elif stop is None:
|
||||
self.stop = lambda attempts, delay: any(f(attempts, delay) for f in stop_funcs)
|
||||
|
||||
else:
|
||||
@ -150,7 +155,10 @@ class Retrying(object):
|
||||
if wait_exponential_multiplier is not None or wait_exponential_max is not None:
|
||||
wait_funcs.append(self.exponential_sleep)
|
||||
|
||||
if wait is None:
|
||||
if wait_func is not None:
|
||||
self.wait = wait_func
|
||||
|
||||
elif wait is None:
|
||||
self.wait = lambda attempts, delay: max(f(attempts, delay) for f in wait_funcs)
|
||||
|
||||
else:
|
||||
|
@ -40,6 +40,13 @@ class TestStopConditions(unittest.TestCase):
|
||||
def test_legacy_explicit_stop_type(self):
|
||||
r = Retrying(stop="stop_after_attempt")
|
||||
|
||||
def test_stop_func(self):
|
||||
r = Retrying(stop_func=lambda attempt, delay: attempt == delay)
|
||||
self.assertFalse(r.stop(1, 3))
|
||||
self.assertFalse(r.stop(100, 99))
|
||||
self.assertTrue(r.stop(101, 101))
|
||||
|
||||
|
||||
class TestWaitConditions(unittest.TestCase):
|
||||
|
||||
def test_no_sleep(self):
|
||||
@ -114,6 +121,13 @@ class TestWaitConditions(unittest.TestCase):
|
||||
def test_legacy_explicit_wait_type(self):
|
||||
r = Retrying(wait="exponential_sleep")
|
||||
|
||||
def test_wait_func(self):
|
||||
r = Retrying(wait_func=lambda attempt, delay: attempt * delay)
|
||||
self.assertEqual(r.wait(1, 5), 5)
|
||||
self.assertEqual(r.wait(2, 11), 22)
|
||||
self.assertEqual(r.wait(10, 100), 1000)
|
||||
|
||||
|
||||
class NoneReturnUntilAfterCount:
|
||||
"""
|
||||
This class holds counter state for invoking a method several times in a row.
|
||||
|
Loading…
Reference in New Issue
Block a user