python mock side_effect或return_value依赖于call_count

时间:2012-11-13 10:19:48

标签: python mocking

为了测试一个轮询函数,我想模拟调用一个子函数,这样第一次调用它就会失败,第二次调用它就会成功。这是一个非常简化的版本:

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 个答案:

答案 0 :(得分:53)

如果我正确理解您的问题,请按setting side_effect to an iterable进行操作。对于你的简单案例:

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

如果您想允许无限数量的通话,请使用itertools cyclechain功能:

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