Unittest for decorator with self arguments

Question

I have a retry decorator that acts based on raised exceptions:

def retry(times, exceptions):
    def decorator(func):
        def wrapped_func(self):
            attempt = 0
            while attempt < times:
                try:
                    return func(self)
                except exceptions:
                    self.logger.info(
                        "Error, retry in process.. attempt {} of {}".format(
                            attempt, times
                        )
                    )
                    attempt += 1
            return func(self)
        return wrapped_func
    return decorator

And will be used in a class function like:

@retry(times=10, exceptions=(ValueError,))
def function_that_fails(self):
    raise ValueError

I'm trying to test this function like:

class sample_class:
    def __init__(self, logger):
        self._logger = logger
        self._count = 0
    
    @property
    def count(self):
        return self._count

    @property
    def logger(self):
        return self._logger
    
    @retry(times=10, exceptions=(ValueError,))
    def sample_func(self):
        self._count += 1
        raise ValueError

class TestUtils(TestCase):
    def setUp(self):
        self._logger = Mock()
        self._sample_class = sample_class(self._logger)
        self._count = 0
    
    def test_retry_function(self):
        with self.assertRaises(ValueError):
            self._sample_class.sample_func()
        self.assertEqual(self._sample_class.count, 10)

    def test_retry_function_2(self):
        @retry(times=10, exceptions=(ValueError,))
        def sample_func():
            self._count += 1
            raise ValueError
        with self.assertRaises(ValueError):
            sample_func()
        self.assertEqual(self._count, 10)

The first test works but the dummy class is something that I would like to avoid, the second one I receive the error:

E   TypeError: wrapped_func() missing 1 required positional argument: 'self'

And it's expected due to the self argument inside the retry function, I tried multiple options with no success, what I'm doing wrong?

user_3799759 · Answer

You could simply use a Mock object for the self parameter instead of creating an entire class:

from unittest.mock import Mock

def test_retry_function_2(self):
    @retry(times=10, exceptions=(ValueError,))
    def sample_func(mock_arg):
        mock_arg._count += 1
        raise ValueError
    mock = Mock(_count=0)
    with self.assertRaises(ValueError):
        sample_func(mock)
    self.assertEqual(mock._count, 10)

Simpler yet, you could use a Mock as the decorated function, and use the built-in call_count in place of a _count that you have to increment yourself:

def test_retry_function_3(self):
    mock = Mock(side_effect=ValueError)
    with self.assertRaises(ValueError):
        retry(10, ValueError)(mock)(Mock())
    self.assertEqual(mock.call_count, 11)

It's even easy to verify that your decorator calls logger.info on each retry:

def test_retry_function_4(self):
    mock_func = Mock(side_effect=ValueError)
    mock_self = Mock()
    with self.assertRaises(ValueError):
        retry(10, ValueError)(mock_func)(mock_self)
    self.assertEqual(mock_func.call_count, 11)
    self.assertEqual(mock_self.logger.info.call_count, 10)

Unittest for decorator with self arguments

1 Answers