How do I mock out file writing to multiple files in python

Question

I'm trying to test a function in which one call results in multiple files being written:

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)

I want my test to check that each call was made to open as expected, and that the content was written:

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"

I've seen an example of handling two files here: Python mock builtin 'open' in a class using two different files

But I can't wrap my head around how to track what is actually written to them.

Answer 1

First, you should never use a mutable object as your default argument to a function, which is an anti-pattern . You should change your function signature to def pull_files(output_files=()) instead.

Then, to your question, you can do a os.chdir to /tmp/ and make a temporary directory, then write files in the temporary folder instead. Don't forget to change your working directory back to what it was after the test.

Another solution is to modify your function slightly so that you are not prepending a prefix ( '/output/' + output_file ). This way, you can pass an io.BytesIO object instead of a path, which will let you modify the contents in-memory.

Answer 2

Here's an example mocking open and returning a StringIO as context:

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"

Similarly you could mock out "regular" use of open (without a context manager).

Edits made: Changed to cover the test cases of the original question.