python mock side_effect or return_value dependent on call_count

Question

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

Answer 1

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'])))