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