python模拟副作用或返回值取决于调用计数

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

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\u count

属性，我应该能够以某种方式使用它

现在我已经通过创建一个自定义的模拟对象解决了这个问题，我用它来monkey patch

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

---

如果我没弄错你的问题，你就按我说的做。对于您的简单案例：

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

如果要允许无限次的调用，请使用

itertools

循环

和

链

功能：

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

---

是的，谢谢，效果很好。然而，值得注意的是，这在

mock

版本

0.7.2

上不起作用，这正是我当时使用的版本。我必须升级到版本

0.8.0

，才能正确使用此功能。