python - 如何模拟查询集以用于 python/Mock 中的 for 循环

Question

我正在编写一些单元测试并想使用 Mock。

给定以下代码:

# the 'real' query set is a Django database model
# qs = SomeDjangoModel.objects.filter(name='some_name')
qs = mock.Mock()
qs.filter.return_value = qs
item = mock.Mock()
item.do_work.return_value = "Some text"
qs.iter.return_value = iter([item])
# below is the code I want to test..
qs = qs.filter(name='some name')
qs = qs.filter(valid_from__lte=Timezone.now())
for obj in qs:
    obj.do_work()

运行时，我得到

TypeError: 'Mock' object is not iterable

我试过打补丁

@mock.patch('__builtin__.iter')

但我似乎无法让它工作。 当 for 循环“使用”查询集时，我还没有成功弄清楚到底发生了什么。

非常感谢您的帮助!

[在第一个解决方案提案之后，使用进一步添加的示例代码进行编辑]

Answer 1

你必须使用 iterator和 MagicMock其中定义了 __iter__

from unittest.mock import Mock, MagicMock
from datetime import datetime

qs = MagicMock()
qs.filter.return_value = qs
item = Mock()
item.do_work.return_value = "Some text"
qs.iterator.return_value = iter([item])
# below is the code I want to test..
qs = qs.filter(name='some name')
qs = qs.filter(valid_from__lte=datetime.now())
for obj in qs:
    obj.do_work()