A python class for retrying flickering unit tests up to a maximum number of times.

retry.py

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