Search code examples
pythonunit-testingpython-mock

Python mock function that returns dictionary multiple times


Background

I have a function my_func that gets a dictionary from another function get_dict and modifies it. On failure, my_func gets retried until it succeeds or called a specified number of times. What complicates this is that the dictionary must have exactly one key-value pair, otherwise other exceptions are raised. Here is a stripped down example of the function:

class MyClass:

    def get_dict(self):
        # Returns a dict

    @retry(3) # Decorator that controls the error handling and retries
    def my_func(self):
        a_dict = self.get_dict()

        if len(a_dict) == 0:
            raise WrongException

        if len(a_dict) > 1:
            raise OtherWrongException

        key, value = a_dict.popitem() # Key-value pair is popped off the dict

        # Do stuff

        raise MyException

Problem

I'm trying to unit test the failure case of my_func, specifically that MyException gets correctly raised after all retries fail. To do this, I mocked the output of get_dict, but because I pop key-value pairs, the original dict is modified. In the real case, get_dict will "refresh" the dict with a new one every time it is called, but the mock dict does not.

This was my unit test attempt:

import MyClass
from unittest import mock, TestCase

class MyTest(TestCase):

    @mock.patch.object(MyClass, "get_dict")
    def test_my_func(self, mock_dict):
        my_class = MyClass()
        mock_dict.return_value = {"key": "value"}

        self.assertRaises(MyException, my_class.my_func)

Instead of the test passing with MyException getting raised, the mocked dictionary return value is modified by the function call and not refreshed for the retry attempts, causing the test to fail when the function raises WrongException.

I have considered setting the side_effect of the mock instead of the return value and just passing a list of dictionaries equal to the number of retry attempts, but that doesn't seem like a good solution.

How can I mock the return value of get_dict so that it returns the unmodified dict every time?


Solution

  • Instead of assigning a mutable object as return value to the mock object:

    mock.patch.object(return_value=...)
    

    You could patch the function with a lambda that returns the mutable object:

    mock.patch.object(new=lambda ...)
    

    Also, the retry decorator should probably be tested independently.

    (retry decorator yanked and adapted from here)

    myclass.py
    def retry(times, exceptions):
        """
        Source: https://stackoverflow.com/a/64030200/20103413
        Retry Decorator
        Retries the wrapped function/method `times` times if the exceptions listed
        in ``exceptions`` are thrown
        :param times: The number of times to repeat the wrapped function/method
        :type times: Int
        :param Exceptions: Lists of exceptions that trigger a retry attempt
        :type Exceptions: Tuple of Exceptions
        """
        def decorator(func):
            def newfn(*args, **kwargs):
                attempt = 1
                while attempt <= times:
                    try:
                        return func(*args, **kwargs)
                    except exceptions as exc:
                        print(
                            '%s thrown when attempting to run %s, attempt '
                            '%d of %d' % (type(exc), func, attempt, times)
                        )
                        attempt += 1
                raise TooManyRetriesException('Maximum attempts exceeded. Terminating retry sequence.')
            return newfn
        return decorator
    
    
    class MissingItemsException(Exception):
        pass
    
    class TooManyItemsException(Exception):
        pass
    
    class TooManyRetriesException(Exception):
        pass
    
    
    class MyClass:
    
        def get_dict(self):
            # Returns a dict
            pass
    
        @retry(3, (RuntimeError,)) # Decorator that controls the error handling and retries
        def my_func(self):
            a_dict = self.get_dict()
    
            if len(a_dict) == 0:
                raise MissingItemsException
    
            if len(a_dict) > 1:
                raise TooManyItemsException
    
            key, value = a_dict.popitem() # Key-value pair is popped off the dict
    
            # Do stuff
    
    test_myclass.py
    from myclass import MyClass, TooManyRetriesException, TooManyItemsException, MissingItemsException
    from unittest import mock, TestCase
    
    class MyTest(TestCase):
    
        @mock.patch.object(MyClass, "get_dict", new=lambda self: {"key": "value"})
        def test_my_func_success(self):
            my_class = MyClass()
    
        @mock.patch.object(MyClass, "get_dict", new=lambda self: {})
        def test_my_func_missing_items(self):
            my_class = MyClass()
            self.assertRaises(MissingItemsException, my_class.my_func)
    
        @mock.patch.object(MyClass, "get_dict", new=lambda self: {"k1": "v1", "k2": "v2"})
        def test_my_func_too_many_items(self):
            my_class = MyClass()
            self.assertRaises(TooManyItemsException, my_class.my_func)