Skip to content

Instantly share code, notes, and snippets.

@jeremystan

jeremystan/retry.py

Last active Sep 14, 2020
Embed
What would you like to do?
A python class for retrying flickering unit tests up to a maximum number of times.
import functools
import logging
class RetryTest:
"""
Decorator for re-trying flickering tests multiple times
to attempt to get them to pass.
Note that it creates a new copy of the test class for each
test run (to ensure no cached data persists between retries).
"""
def __init__(self, max_num_tries=3):
self.max_num_tries = max_num_tries
def __get__(self, obj, objtype):
"""Support instance methods."""
return functools.partial(self.__call__, obj)
def __call__(self, func, *args, **kwargs):
@functools.wraps(func)
def wrapper(*args, **kwargs):
num_tries = 0
passed = False
exception = None
while not passed and num_tries < self.max_num_tries:
try:
if num_tries == 0:
# Use the original unittest object
result = func(*args, **kwargs)
else:
logging.debug(f"test try {num_tries} failed, retrying")
# Re setup the class and create a new instance to clear any cached data
obj = args[0]
_cls = obj.__class__
obj.tearDown()
_cls.tearDownClass()
_cls.setUpClass()
new_obj = args[0].__class__()
new_obj.setUp()
# Re-run the test
result = func(new_obj, *args[1:], **kwargs)
except AssertionError as e:
exception = e
else:
passed = True
finally:
num_tries += 1
if not passed:
raise exception
else:
return result
return wrapper
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment