Testing celery task retries

Had the following testing scenario the other day. Celery task with a possible exception that may be raised. When raised, the task should should run again but with a different kwarg.

Example function to test
@task()
def foobar(flag=False):  
    try:
        do_something()
    except SpecificException as e:
        log.debug('Something specific happend')
        raise retry(exc=ex, kwargs={'flag': True})
Test
from celery.exceptions import Retry


class TaskTestCase(unittest.TestCase)  
    def test_retry(self):
        mock_retry = mock.patch('foobar.retry', autospec=True).start()
        mock_retry.side_effect = Retry
        mock_do_something = mock.patch('do_something', autospec=True).start()

        e = Exception('Whoops')
        mock_do_something.side_effect = [e, None]

        with self.assertRaises(Retry):
            foobar().apply()

        self.mock_retry.assert_called_with(exc=e, kwargs={'flag': True})

Lets break this down....

Your celery task should have a retry function. Mock this. Give it a side effect of celery.exceptions.Retry. This sideeffect is expected and task will fail if run without it.

mock_retry = mock.patch('foobar.retry', autospec=True).start()  
mock_retry.side_effect = Retry  

The called function should fail once and then run successfully the second time. Python mock allows patched objects to use a list of side effects.

mock_do_something.side_effect = [e, None]  

After running the function to test, the retry that we patched can now be tested for args. In this specific case the retry function is only called once.

self.mock_retry.assert_called_with(exc=e, kwargs={'flag': True})  

I would be curious to know how often this type of task pattern is used, if ever. It seems to have a place so I hope this testing pattern is useful.