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?
CodePudding user response：
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)
CodePudding user response：
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}