python mock的side_effect或return_value取决于call_count

为了测试一个轮询函数，我想模拟调用一个子函数，使得第一次调用时会失败，而第二次调用时会成功。这是一个非常简化的版本：

poll_function(var1):
    value = sub_function(var1)  # 第一次调用将返回None
    while not value:
        time.sleep(POLLING_INTERVAL)  
        value = sub_function(var1) # 后续调用将返回一个字符串，例如"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())