python retry with tenacity, disable `wait` for unittest

Question

I am using the tenacity library to use its @retry decorator.

I am using this to make a function which makes a HTTP-request "repeat" multiple times in case of failure.

Here is a simple code snippet:

@retry(stop=stop_after_attempt(7), wait=wait_random_exponential(multiplier=1, max=60))
def func():
   ...
   requests.post(...)

The function uses the tenacity wait-argument to wait some time in between calls.

The function together with the @retry-decorator seems to work fine.

But I also have a unit-test which checks that the function gets called indeed 7 times in case of a failure. This test takes a lot of time because of this wait in beetween tries.

Is it possible to somehow disable the wait-time only in the unit-test?

Answer 1

You can use unittest.mock module to mock some elements of tentacity library. In your case all decorators you use are classes e.g. retry is a decorator class defined here. So it might be little bit tricky, but I think trying to

mock.patch('tentacity.wait.wait_random_exponential.__call__', ...)

may help.

Answer 2

The solution came from the maintainer of tenacity himself in this Github issue: https://github.com/jd/tenacity/issues/106

You can simply change the wait function temporarily for your unit test:

from tenacity import wait_none

func.retry.wait = wait_none()

Answer 3

Thanks to discussion here, I found an elegant way to do this based on code from @steveb:

from tenacity import retry, stop_after_attempt, wait_exponential


@retry(reraise=True, stop=stop_after_attempt(5), wait=wait_exponential(multiplier=1, min=4, max=10))
def do_something_flaky(succeed):
    print('Doing something flaky')
    if not succeed:
        print('Failed!')
        raise Exception('Failed!')

And tests:

from unittest import TestCase, mock, skip
from main import do_something_flaky


class TestFlakyRetry(TestCase):
    def test_succeeds_instantly(self):
        try:
            do_something_flaky(True)
        except Exception:
            self.fail('Flaky function should not have failed.')

    def test_raises_exception_immediately_with_direct_mocking(self):
        do_something_flaky.retry.sleep = mock.Mock()
        with self.assertRaises(Exception):
            do_something_flaky(False)

    def test_raises_exception_immediately_with_indirect_mocking(self):
        with mock.patch('main.do_something_flaky.retry.sleep'):
            with self.assertRaises(Exception):
                do_something_flaky(False)

    @skip('Takes way too long to run!')
    def test_raises_exception_after_full_retry_period(self):
        with self.assertRaises(Exception):
            do_something_flaky(False)

Answer 4

mock the base class wait func with:

mock.patch('tenacity.BaseRetrying.wait', side_effect=lambda *args, **kwargs: 0)

it always not wait

Answer 5

After reading the thread in tenacity repo (thanks @DanEEStar for starting it!), I came up with the following code:

@retry(
    stop=stop_after_delay(20.0),
    wait=wait_incrementing(
        start=0,
        increment=0.25,
    ),
    retry=retry_if_exception_type(SomeExpectedException),
    reraise=True,
)
def func() -> None:
    raise SomeExpectedException()


def test_func_should_retry(monkeypatch: MonkeyPatch) -> None:
    # Use monkeypatch to patch retry behavior.
    # It will automatically revert patches when test finishes.
    # Also, it doesn't create nested blocks as `unittest.mock.patch` does.

    # Originally, it was `stop_after_delay` but the test could be
    # unreasonably slow this way. After all, I don't care so much
    # about which policy is applied exactly in this test.
    monkeypatch.setattr(
        func.retry, "stop", stop_after_attempt(3)
    )

    # Disable pauses between retries.
    monkeypatch.setattr(func.retry, "wait", wait_none())

    with pytest.raises(SomeExpectedException):
        func()

    # Ensure that there were retries.
    stats: Dict[str, Any] = func.retry.statistics
    assert "attempt_number" in stats
    assert stats["attempt_number"] == 3

I use pytest-specific features in this test. Probably, it could be useful as an example for someone, at least for future me.

Answer 6

I wanted to override the retry function of the retry attribute and while that sounds obvious, if you are playing with this for the first time it doesn't look right but it is.

sut.my_func.retry.retry = retry_if_not_result(lambda x: True)

Thanks to the others for pointing me in the right direction.

Answer 7

You can mock tenacity.nap.time in conftest.py in the root folder of unit test.

@pytest.fixture(autouse=True)
def tenacity_wait(mocker):
    mocker.patch('tenacity.nap.time')