为了测试一个轮询函数,我想模拟调用一个子函数,这样第一次调用它就会失败,第二次调用它就会成功。这是一个非常简化的版本:
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())
答案 0 :(得分:53)
如果我正确理解您的问题,请按setting side_effect
to an iterable进行操作。对于你的简单案例:
>>> mock_poll = Mock(side_effect=[None, 'data'])
>>> mock_poll()
None
>>> mock_poll()
'data'
如果您想允许无限数量的通话,请使用itertools
cycle
和chain
功能:
>>> mock_poll = Mock(side_effect=chain(['first'], cycle(['others'])))