Python mock side_effect или return_value зависит от call_count

Чтобы проверить функцию опроса, я хочу высмеять вызов подфункции так, что в первый раз, когда она будет вызвана, она потерпит неудачу, а второй раз она будет вызвана. Здесь очень упрощенная версия:

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

Возможно ли это сделать с объектом Mock из фреймворка Mock? Я знаю, что объекты Mock имеют атрибут call_count, который я должен каким-то образом использовать.

Сейчас я решил это, создав собственный макет-объект, который я использую для патча обезьяны sub_function(), но я считаю, что должен быть более менее подробный способ сделать это:

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

Ответ 1