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

最新推荐文章于 2025-06-11 14:51:15 发布

原创最新推荐文章于 2025-06-11 14:51:15 发布 · 3.2k 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #单元测试

Python 同时被 3 个专栏收录

25 篇文章

订阅专栏

Python拾遗

24 篇文章

订阅专栏

单元测试

2 篇文章

订阅专栏

本文探讨了Python中返回生成器函数的内存效率，并展示了在单元测试中如何使用mock模拟这类函数的返回值。通过示例代码解释了如何使用unittest.mock库中的Mock和MagicMock来替代并控制生成器的行为。

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())

以上代码的执行结果如下：