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 << ----------------------

Bacitracin

Well, you did mock do_something(). Mocking a method doesn't then call the original, with our without side effects. – Manning

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.

Manning

What should we do in case if do_something is a class method. It can't be so easy saved, due to since it is a class method we somehow should be able to pass self. Any ideas? btw since it is some internal logic it couldn't be just instantiated in test suite – Upthrow

@xiº: why do you need to have access to do_something at all? It is not part of the code-under-test, so from a unit-testing point of view you should just mock it. Any reference to a callable can be saved however, it doesn't matter if it is a class method. – Manning

I need something like side_effect = [original_valid_call(), original_valid_call(), ValueError]. But can't manage how to pass an instance in valid call. – Upthrow

In my case do_smth is just underlying layer. – Upthrow

@xiº: why not provide mock objects instead? Otherwise, how would you normally call that class method? – Manning

I just need original behavior of that method for several times and exception on nth step. It is method of class of internal library, my actually tested code just bothering about either there are any exception or not. – Upthrow

@xiº: perhaps you can post a new question on that then? There is nothing special about calling a class method, and I'm not certain what your issue is. Include a minimal reproducible example to demonstrate where you are stuck. – Manning

here it is #43191642 – Upthrow

Answer 2

In case someone just needs to return values after the stopIterations happens, you just need to return the iterator, not the exception.

responses = [1,2,3,4,5]

def requests_side_effect(*args, **kwargs):
    try:
        return next(responses)
    except StopIteration:
        # raised all exceptions, call original
        return default_request_mock

This will retrieve values for the first 5 calls and then return the default value.

Bertrando