Python helper functions for implementing retry with exponential backoff and custom error checking

retry_utils.py

class RetryUtils():
    """
        Provides helper functions for implementing retry with exponential backoff.
    """
    
    def retry_with_backoff(
            function: Callable,
            function_args: list[Any],
            allowed_exceptions: list[str] = [],
            retries: int=5, 
            backoff_in_seconds:int=1,
            mock_retry: bool=False
        ) -> Any:
        """
            Helper function that wraps a function in retry with exponential backoff.

            Attributes
            ----------
            function                : Callable
                the function to be wrapped in the retry with exponential backoff.
            function_args           : list[Any]
                the arguments to be passed to the function.
            allowed_exceptions      : list[str]
                list of exception names that are allowed and will permit retry.
                If this list is empty, any exception will be permitted.
            retries                 : int
                the number of retries to be attempted.
            backoff_in_seconds      : int
                the base number of seconds (upon which exponential backoff will be generated)
                to wait between retry attempts.
            mock_retry              : bool
                if True, the function does not sleep. Useful for use in mocks or unit tests.

            Returns
            ----------
            function_result         : Any
                the return result of the Callable function
        """

        count = 0

        while True:

            try:

                return function(*function_args)
            
            except Exception as ex:

                ex_type = type(ex).__name__
                ex_msg = str(ex)

                if len(allowed_exceptions) > 0:

                    exception_match = False

                    for allowed_exception in allowed_exceptions:

                        if (
                                allowed_exception.lower() in ex_type.lower()
                                or allowed_exception.lower() in ex_msg.lower()
                        ):

                            exception_match = True
                            
                            break

                    if (exception_match):

                        info_msg = (
                            f"An allowed exception {ex_type}, {ex_msg} has occurred during execution for " +
                            f"function: {str(function)} with function_args: {function_args}. As this is an allowed " +
                            "exception, retry with exponential backoff will be attempted."
                        )

                        logger.info(info_msg)

                    else:

                        error_msg = (
                            f"An unexpected exception {ex_type}, {ex_msg} has occurred during execution for " +
                            f"function: {str(function)} with function_args: {function_args}. This exception " +
                            f"does not match one of the provided allowed_exceptions: {', '.join(allowed_exceptions)}. " +
                            "Raising error and not attempting retry."
                        )

                        logger.error(error_msg)

                        raise ValueError(ex)

                else:

                    info_msg = (
                        f"An unexpected exception {ex_type}, {ex_msg} has occurred during execution for " +
                        f"function: {str(function)} with function_args: {function_args}. As no explicit allowed " +
                        "exceptions have been defined, retry with exponential backoff will be attempted."
                    )

                    logger.warning(info_msg)


                if count == retries:

                    raise ValueError(
                        f"Maximum number of retries: {retries} exceeded for function {str(function)} " +
                        f"with function_args: {function_args}."
                    )
                
                if not mock_retry:

                    sleep = (backoff_in_seconds * 2**count + random.uniform(0, 1))

                    time.sleep(sleep)

                count += 1

test_retry_utils.py

import RetrtyUtils

import pytest


def test_retry_with_backoff_explicit_allowed_exceptions():

    def test_function():

        raise LookupError("Allowed exception for test case")


    with pytest.raises(ValueError) as e_info:
        RetryUtils.retry_with_backoff(
            function=test_function,
            function_args=[],
            allowed_exceptions=["LookupError"],
            retries=1,
            backoff_in_seconds=1,
            mock_retry=True
        )

        assert 'Maximum number of retries: 1 exceeded for function' in str(e_info.value)


def test_retry_with_backoff_implicit_allowed_exceptions():

    def test_function():

        raise LookupError("Allowed exception for test case")


    with pytest.raises(ValueError) as e_info:
        RetryUtils.retry_with_backoff(
            function=test_function,
            function_args=[],
            retries=1,
            backoff_in_seconds=1,
            mock_retry=True
        )

        assert 'Maximum number of retries: 1 exceeded for function' in str(e_info.value)


def test_retry_with_backoff_success_no_args():

    test_string = "Hello from test_retry_with_backoff_success test case"

    def test_function():

        return test_string

    result = RetryUtils.retry_with_backoff(
        function=test_function,
        function_args=[],
        retries=1,
        backoff_in_seconds=1,
        mock_retry=True
    )

    assert result
    
    assert result == test_string


def test_retry_with_backoff_success_with_args():

    test_string = "argument1,55,argument3"

    def test_function(arg1, arg2, arg3):

        return ','.join(str(e) for e in [arg1, arg2, arg3])

    result = RetryUtils.retry_with_backoff(
        function=test_function,
        function_args=["argument1", 55, "argument3"],
        retries=1,
        backoff_in_seconds=1,
        mock_retry=True
    )

    assert result
    
    assert result == test_string