python mock side_effect or return_value dependent on call_count
- by user18380
To test a polling function I want to mock the calling of a sub function so that the first time it is called it will fail, and the second time it is called it will succeed. Here's a very simplified version of it:
poll_function(var1):
value = sub_function(var1) # First call will return None
while not value:
time.sleep(POLLING_INTERVAL)
value = sub_function(var1) # A subsequent call will return a string, e.g "data"
return value
Is this possible to do with a Mock object from the mock framework? I know Mock objects have a call_count attribute I should be able to use somehow.
Right now I've solved it by creating a custom mock object that I use to monkey patch sub_function(), but I feel there should be a better less verbose way of doing it:
def test_poll():
class MyMock(object):
def __init__(self, *args):
self.call_count = 0
def sub_function(self, *args, **kwargs):
if self.call_count > 1:
return "data"
else:
self.call_count += 1
return None
my_mock = MyMock()
with patch('sub_function', my_mock.sub_function):
ok_(poll_function())
