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?

Asked By: btj29

||

Source

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)

Answered By: Fractalism

Python mock function that returns dictionary multiple times

Question:

Background

Problem

Answers:

myclass.py

test_myclass.py