Python函数可以通过返回generator
的方法节约内存。所谓的返回generator
,就是通过yield
语句返回一系列的值。这些值只有在程序的返回结果被迭代(iterate)的情况下才会被真正返回。有关yield
语句的工作原理和generator
的详细解释在这篇文章中。在此不做详细讨论。
当我们使用yield
来返回一个方法的执行结果时,实际上返回的是一个generator
。而这个generator
中的值只有在之后循环这个generator
(例如使用for
循环)时才会被真正返回。这个循环的过程实际上调用的是genereator
的__iter__
函数。所以我们只要mock
所返回的generator
的__iter__
函数的返回值即可。可以使用MagicMock()
或者mock.patch()
来创建mock
。
如下面的例子所示:
例如我们有如下类需要测试:
class example(object):
def func(self):
yield from [1, 2, 3]
eg = example()
list(eg.func())
以上代码的执行结果如下:
[1, 2, 3]
那么我们mock
该类及其func
函数的方法如下,注意为了配置generator
的返回值,我们实际上mock
了func
函数所返回的generator
的__iter__()
的返回值。
from mock import MagicMock
mock_example = MagicMock()
mock_example.func.return_value.__iter__.return_value = iter([1, 2, 3])
list(mock_example.func())
以上代码的执行结果如下:
[1, 2, 3]
可见我们成功的mock
了一个返回generator
的函数。