python - 如何使用 Mock 框架更改模拟方法的输出？

Question

我有一个像下面这样的模块:

'''example.py'''
from some.module import Session

def init_system():
    var_one = Session.query(SQLAModelA).all()
    var_two = [(x.attr_B_a, x) for x in Session.query(SQLAModelB).all()]

    for value in var_one:
        something = value.attr_A_a + value.attr_A_c
    # Will do something with var_two etc (which has different attributes)
    something_else = [x.attr_B * 2 for x in var_two]
    return result

session 是指与数据库对话的 SQLAlchemy session 处理程序。出于测试目的，我需要能够模拟 Session 对象(或者可能只是 Session 对象？)中的查询方法，这样我就可以在 init_system 实际与数据库交互的情况下进行测试。我如何使用 Mock framework做这个？

真正让我困惑的一点是，如果我模拟 query() 方法，我如何根据传递给 query() 的类更改它的所有方法的输出??

我希望我可以按照以下方式做一些事情:

def test_init_numbers():
    import
    my_mock_object = MagicMock()
    with patch('some.module.Session.query', my_mock_object):
        result = example.init_system()
        assert result == ['expected', 'result']

但是，我不知道我需要对 my_mock_object 做些什么才能使其正确模拟 query()。我想它需要实现返回迭代器的 all() 方法。而且，我知道我可以创建更多模拟对象以放入返回的迭代器中，确保它们具有适当属性的值，例如:

var_one_element_mock = MagicMock(spec=SQLAModelA)
var_one_element_mock.attr_A_a = 12

var_two_element_mock = MagicMock(spec=SQLAModelB)
var_one_element_mock.attr_B = 100

当然，我可以创建那些 MagicMock 对象的不同版本的迭代器，这些对象可以被 init_system 函数使用，例如，在 for value in var_one 中行。

但我不知道如何将所有模拟对象拼凑在一起，所以我只能模拟 Session.query 方法，同时考虑到传递给它的类会影响内容它返回的可迭代对象。

谢谢!

Answer 1

我不知道您是否仍需要这方面的帮助，但如果您想控制调用 MagicMock 对象时发生的情况，您可以传递一个函数作为其 side_effect 参数。 这样的事情应该有效。

def test_init_numbers():
    import
    my_mock_object = MagicMock()
    def return_iterator():
      ## create and return your iterator here

    my_mock_object.all = MagicMock(side_effect=return_iterator)
    with patch('some.module.Session.query', my_mock_object):
        result = example.init_system()
        assert result == ['expected', 'result']