Skip to content

Instantly share code, notes, and snippets.

@martin056
Last active August 13, 2019 17:06
Show Gist options
  • Save martin056/1600c20bbc1dac27b9543a8eb9ee5668 to your computer and use it in GitHub Desktop.
Save martin056/1600c20bbc1dac27b9543a8eb9ee5668 to your computer and use it in GitHub Desktop.
This is an example of a way to have a unittest mock that returns different values each time it's called
# in example.py
def inner_func(x=None):
if x is None:
return 1
return x + 1
def big_func():
a = inner_func()
b = inner_func(a)
c = inner_func(b)
return c
# in mocks.py
from unittest.mock import Mock
class MyMock(Mock):
def __call__(self, *args, **kwargs):
if isinstance(self.return_values, MyMock):
"""
By default, when you try to access a property of a Mock which is not set
you will get a Mock object from the same type (in this case MyMock).
Because of the way unittest treats our mocks, we cannot use neither the self.__class__
nor type(self) in the if-condition.
"""
raise ValueError('Provide `return_values` to the mock')
super().__call__(*args, **kwargs)
call = self.call_count - 1
if (len(self.return_values) >= call):
return self.return_values[-1]
return self.return_values[call]
# in tests.py
from unittest.mock import patch
from unittest import TestCase
from .example import big_func
from .mocks import MyMock
class BigFuncTests(TestCase):
@patch('<path_to_inner_func>.inner_func', new_callable=MyMock)
def test_mega_shano_use_case(self, inner_mock):
inner_mock.return_values = [1, 2, 10]
res = big_func()
self.assertEqual(res, 10)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment