我如何模拟文件写入 python 中的多个文件

Question

我正在尝试测试一个函数，其中一次调用会导致写入多个文件：

def pull_files(output_files=[]):
    for output_file in output_files:
        content = get_content_from_server(output_file)
        with open('/output/' + output_file, "wb") as code:
            code.write(content)

我希望我的测试检查每个调用是否按预期打开，以及内容是否已写入：

def test_case(self):
    pull_files("file1.txt", "file2.txt")

    # Assert open("file1.txt", "wb") was called
    # Assert "file 1 content" was written to "file1.txt"

    # Assert open("file2.txt", "wb") was called
    # Assert "file 2 content" was written to "file2.txt"

我在这里看到了一个处理两个文件的例子： Python mock builtin 'open' in a class using two different files

但是我无法理解如何跟踪实际写入它们的内容。

Answer 1

首先，永远不要使用可变对象作为函数的默认参数，这是一个反模式。 您应该将函数签名更改为def pull_files(output_files=()) 。

然后，对于您的问题，您可以对/tmp/执行os.chdir并创建一个临时目录，然后将文件写入临时文件夹中。 不要忘记将您的工作目录改回测试后的状态。

另一种解决方案是稍微修改你的函数，这样你就不会在前面加上前缀（ '/output/' + output_file ）。 这样，您可以传递一个io.BytesIO对象而不是路径，这将允许您修改内存中的内容。

Answer 2

这是一个StringIO open并返回StringIO作为上下文的示例：

from io import StringIO

def my_function(*fns):
    for i, fn in enumerate(fns):
        with open(fn, "wt") as fp:
            fp.write("content %d" % i)


string_io_one = StringIO()
string_io_two = StringIO()
with mock.patch("%s.open" % __name__) as open_mock:
    open_mock.return_value.__enter__.side_effect = [string_io_one, string_io_two]
    my_function("file1.txt", "file2.txt")

    assert open_mock.called_with("file1.txt")
    string_io_one.seek(0)
    assert string_io_one.read() == "content 0"
    assert open_mock.called_with("file2.txt")
    string_io_two.seek(0)
    assert string_io_two.read() == "content 1"

同样，您可以模拟open “常规”使用（没有上下文管理器）。

所做的编辑：更改以涵盖原始问题的测试用例。