python mock side_effect or return_value dependent on call_count

mocking, python

Solution

If I understand your question correctly, you do it by setting `side_effect` to an iterable. For your simple case:

>>> mock_poll = Mock(side_effect=[None, 'data'])
>>> mock_poll()
None
>>> mock_poll()
'data'

If you want to allow for an unlimited number of calls, use the `itertools` `cycle` and `chain` functions:

>>> mock_poll = Mock(side_effect=chain(['first'], cycle(['others'])))

Problem

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()) ```

Original source