Python单元测试中mock一个返回generator的函数

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的函数。