Mock side effect only X number of times

Question

I have a celery retry task that I would like to test that it retries until successful. Using mock's side_effect, I can fail it for a set number of executions and then passing None, clear the side effect. However, the method the task is calling doesn't execute at that point, it just doesn't have an exception. Is there a way to clear the side effect, and still have the method being mocked execute as normal?

I can test that it is called 'x' number of times (ie. repeat until successful) and then in a separate test, assert it does what is supposed to, but was wondering if there was a way to do both in one test.

tasks.py:

import celery

@celery.task(max_retries=None)
def task():
    print "HERE"
    try:
        do_something("TASK")
    except Exception as exc:
        print exc
        raise task.retry(exc=exc)

def do_something(msg):
    print msg

Test:

import ....

class TaskTests(test.TestCase):

    @mock.patch('tasks.do_something')
    def test_will_retry_until_successful(self, action):
        action.side_effect = [Exception("First"), Exception("Second"), Exception("Third"), None]
        tasks.task.delay()
        self.assert.... [stuff about task]

Results: fails three times and then "succeeds" but do_something() never prints. action.call_count equals 4. I would like to see that the blank line following the last 'HERE' would be print of 'TASK'.

-------------------- >> begin captured stdout << ---------------------
HERE
First
HERE
Second
HERE
Third
HERE

--------------------- >> end captured stdout << ----------------------

Answer 1

You mocked do_something(). A mock replaces the original entirely; your choices are to either have the side effect (raise or return a value from the iterable) or to have the normal mock operations apply (returning a new mock object).

In addition, adding None to the side_effect sequence doesn't reset the side effect, it merely instructs the mock to return the value None instead. You could add in mock.DEFAULT instead; in that case the normal mock actions apply (as if the mock had been called without a side effect):

@mock.patch('tasks.do_something')
def test_will_retry_until_successful(self, action):
    action.side_effect = [Exception("First"), Exception("Second"), Exception("Third"), mock.DEFAULT]
    tasks.task.delay()
    self.assert.... [stuff about task]

If you feel your test must end with calling the original, you'll have to store a reference to the original, unpatched function, then set the side_effect to a callable that will turn around and call the original when the time comes:

# reference to original, global to the test module that won't be patched
from tasks import do_something

class TaskTests(test.TestCase):
    @mock.patch('tasks.do_something')
    def test_will_retry_until_successful(self, action):
        exceptions = iter([Exception("First"), Exception("Second"), Exception("Third")])
        def side_effect(*args, **kwargs):
            try:
                raise next(exceptions)
            except StopIteration:
                # raised all exceptions, call original
                return do_something(*args, **kwargs)
        action.side_effect = side_effect
        tasks.task.delay()
        self.assert.... [stuff about task]

I cannot, however, foresee a unittesting scenario where you'd want to do that. do_something() is not part of the Celery task being tested, it is an external unit, so you should normally only test if it was called correctly (with the right arguments), and the correct number of times.