retry utility for Python

retry.py

import time
import inspect
import types
 
def _assert(func):
    if not func():
        if isinstance(func, types.FunctionType):
            raise AssertionError(inspect.getsource(func).strip())
        else:
            raise AssertionError()
 
def wait(func, *args, **kargs):
    retry(lambda: _assert(func), *args, **kargs)
 
def retry(func, attempts=0, delay=0, timeout=0):
    start = time.time()
    count = 0
 
    # do-while
    while True:
        try:
            count += 1
            return func()
        except Exception as e:
            last = e
 
        if timeout and time.time() - start + delay >= timeout:
            raise last
        if attempts and count >= attempts:
            raise last
 
        time.sleep(delay)

test_retry.py

import pytest
import time
from mock import Mock
 
from utils.retry import *
 
def test_retry_retval(value=1234):
    """`retry` should return the value from `func` when successful."""
    mock = Mock(return_value=value)
 
    retval = retry(mock)
    
    assert retval == value
    assert mock.call_count == 1
 
def test_retry_eval():
    """`retry` should re-evaluate `func` when it throws an exception."""
    mock = Mock()
    mock.side_effect = [Exception('Boom!'), 1234]
 
    retval = retry(mock)
 
    assert retval == 1234
    assert mock.call_count == 2
 
@pytest.mark.parametrize("a", [1, 2, 10])
def test_retry_attempts(a):
    """`retry` should re-evaluate `func` up to `attempts` times."""
    mock = Mock()
    mock.side_effect = Exception('Boom!')
 
    with pytest.raises(Exception) as e:
        retry(mock, attempts=a)
    
    assert str(e.value) == 'Boom!'
    assert mock.call_count == a
 
@pytest.mark.parametrize("t", [1, 2, 10])
def test_retry_timeout(t, d=0.1):
    """`retry` should re-evaluate `func` until `timeout` is reached."""
    mock = Mock()
    mock.side_effect = Exception('Boom!')
 
    start = time.time()
    with pytest.raises(Exception) as e:
        # delay needs to be non-zero for mocktime
        retry(mock, delay=d, timeout=t)
 
    assert str(e.value) == 'Boom!'
    assert time.time() - start >= t - d
 
@pytest.mark.parametrize("d", [0.1, 0.2, 1.0])
def test_retry_delay(d):
    """`retry` should re-evaluate `func` with `delay` sec between attempts."""
    mock = Mock()
    mock.side_effect = Exception('Boom!')
 
    with pytest.raises(Exception) as e:
        retry(mock, delay=d, timeout=1)
 
    assert str(e.value) == 'Boom!'
    assert mock.call_count <= 1/d + 1
 
def test_wait_nowait():
    """`wait` should not delay if `func` returns True immediately."""
    mock = Mock(return_value=True)
 
    start = time.time()
    wait(mock)    
 
    assert time.time() == start
 
@pytest.mark.parametrize("a", [1, 2, 10])
def test_wait_attempts(a):
    """`wait` should re-evaluate `func` up to `attempts` times."""
    mock = Mock(return_value=False)
 
    with pytest.raises(AssertionError):
        wait(mock, attempts=a)
 
    assert mock.call_count == a
 
@pytest.mark.parametrize("t", [1, 2, 10])
def test_wait_timeout(t, d=0.1):
    """`wait` should raise an exception after `timeout`."""
    mock = Mock(return_value=False)
 
    start = time.time()
    with pytest.raises(AssertionError):
        wait(mock, delay=d, timeout=t)
 
    assert time.time() - start >= t - d
    
@pytest.mark.parametrize("d", [0.1, 0.2, 1.0])
def test_wait_delay(d):
    """`wait` should re-evaluate `func` with `delay` sec between attempts."""
    mock = Mock(return_value=False)
 
    with pytest.raises(AssertionError):
        wait(mock, delay=d, timeout=1)
 
    assert mock.call_count <= 1/d + 1