Python mock function that returns dictionary multiple times

Question

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?

Answer 1

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)

Answer 2

This is a simplified example of how I would in general mock-up a method so that on successive calls it returns different values. You would, of course, modify mocked_get_dict to return whatever you needed and my_func to test approproately.

from unittest import mock

def retry(cnt):
    def retry_internal(f):
        def wrapper(*args, **kwargs):
            nonlocal cnt
            while cnt:
                try:
                    result = f(*args, **kwargs)
                except Exception as e:
                    print(e, end='')
                    cnt -= 1
                    if (cnt):
                        print(' retrying...')
                    else:
                        print()
                else:
                    return result
        return wrapper
    return retry_internal

class MyClass:

    def get_dict(self):
        # Returns a dict
        return {}

    @retry(3)
    def my_func(self):
        d = self.get_dict()
        if not 'x' in d:
            raise ValueError(f'"x" key is missing: {d}')
        print(d)

mock_call_number = -1
mock_return_values = [
    {'a': 1},
    {'b': 2},
    {'x': 3}
]

def mocked_get_dict(self):
    global mock_call_number

    mock_call_number += 1
    if mock_call_number >= len(mock_return_values):
        raise 'Not enough mocked values.'
    return mock_return_values[mock_call_number]

class MyTest:
    @mock.patch.object(MyClass, 'get_dict', mocked_get_dict)
    def test_my_func(self):
        my_class = MyClass()
        my_class.my_func()

MyTest().test_my_func()

Prints:

"x" key is missing: {'a': 1} retrying...
"x" key is missing: {'b': 2} retrying...
{'x': 3}