Python: mock reading and writing files

20 April 2024

Since Python 3.3, the standard library has included unitest.mock, which provides utilities for mocking out parts of the system for testing.

These notes cover how to use unittest.mock when reading from one or more files, or writing to a single file.

Mock reading a single file §

Using the patch context manager, we mock out all calls to open(...) (i.e. "builtins.open") within that block.

By using the special mock object mock_open(read_data=...), we can specify what data will be returned when readlines() is called.
By keeping a reference to the mock object (m), we can query which arguments were used in the call to open(...).

As an example:

from unittest.mock import patch, mock_open

DATA = """\
line 1
line 2
line 3
"""

def read_from_file(filename: str):
    with open(filename, "r") as handle:
        return handle.readlines()

def test_read_from_file():
    m = mock_open(read_data=DATA)
    with patch("builtins.open", m):
        lines = read_from_file("my_file.txt")

    m.assert_called_once_with("my_file.txt", "r")

    assert lines[0] == "line 1\n"
    assert lines[1] == "line 2\n"
    assert lines[2] == "line 3\n"

These examples are written for running under pytest, but the body of the test function should work just the same regardless of test framework. No pytest specific functionality is used for the mocking, only the Python standard library.

Mock reading multiple files §

A limitation of the previous example is that all calls to open(...) within that block are mocked out with the same mock object. This is fine if you are only mocking one call to open(...), but if you have multiple calls, you need each one to be assigned a different mock_open() object.

To support mocking multiple different calls to open(...), we need to supply a separate mock_open() object depending upon which filename is being requested:

First, create a filename_to_mock_open() function. This takes a filename and returns a mock_open(...)(). Note the additional parentheses.
Next, create a mock_open_multiple(on_open=...) function. It takes one argument (on_open), which is the function to execute instead of open().

Putting it all together, we get this:

from unittest.mock import patch, mock_open

DATA_FILE_1 = """\
file 1, line 1
file 1, line 2
file 1, line 3
"""

DATA_FILE_2 = """\
file 2, line 1
file 2, line 2
file 2, line 3
"""

def filename_to_mock_open(filename: str, *args, **kwargs):
    """Returns a `mock_open` object for the corresponding filename.

    :param filename: Name of file to open.
    :param args: Positional arguments passed into `open`, e.g. "r".
    :param kwargs: Keyword arguments passed into `open`.
    """
    return {
        "my_file_1.txt": mock_open(read_data=DATA_FILE_1)(),
        "my_file_2.txt": mock_open(read_data=DATA_FILE_2)(),
    }[filename]

def mock_open_multiple(on_open):
    """A nested `mock_open` which permits using multiple different `mock_open`
    objects depending upon the given parameters (e.g. filename). This allows
    us to patch some code that calls `open()` multiple times, and handle each
    call to `open()` differently.

    :param on_open: Function to call instead of `open`. The given function
        gets passed all the parameters that `open` would have received.
    """
    mock_files = mock_open()
    mock_files.side_effect = on_open
    mock_files.return_value = None
    return mock_files

def read_from_files(filenames: list[str]):
    ret = {}
    for filename in filenames:
        with open(filename, "r") as handle:
            ret[filename] = handle.readlines()
    return ret

def test_read_from_multiple_files():
    m = mock_open_multiple(on_open=filename_to_mock_open)
    with patch("builtins.open", m):
        lines = read_from_files(filenames=["my_file_1.txt", "my_file_2.txt"])

    assert lines["my_file_1.txt"][0] == "file 1, line 1\n"
    assert lines["my_file_1.txt"][1] == "file 1, line 2\n"
    assert lines["my_file_1.txt"][2] == "file 1, line 3\n"

    assert lines["my_file_2.txt"][0] == "file 2, line 1\n"
    assert lines["my_file_2.txt"][1] == "file 2, line 2\n"
    assert lines["my_file_2.txt"][2] == "file 2, line 3\n"

You may want to organise things differently:

The mock_open_multiple() function is generic, and could be added to a set of common utility functions for testing.
The filename_to_mock_open() function is specific to code under test.

Mock writing a single file §

We can mock writing to a file by using patch to replace open() with a mock object (m), and then we can assert that the mock object is called as expected:

from unittest.mock import patch, call

def write_to_file(filename: str):
    with open(filename, "w") as handle:
        handle.write("line 1\n")
        handle.write("line 2\n")
        handle.write("line 3\n")

def test_write_to_file():
    with patch("builtins.open") as m:
        write_to_file("my_file.txt")

    m.assert_has_calls(
        [
            call("my_file.txt", "w"),
            call().__enter__(),
            call().__enter__().write("line 1\n"),
            call().__enter__().write("line 2\n"),
            call().__enter__().write("line 3\n"),
            call().__exit__(None, None, None),
        ]
    )

Note that, because we use a context manager for writing (e.g. with open(filename, "w")), all the calls to write occur within that context (e.g. call().__enter__().write(...)).

If the order does not matter, use assert_has_calls(..., any_order=True).

Conclusion §

Python makes it easy to mock out files for testing, but it’s even easier if you don’t need to mock in the first place.

As a rule of thumb, keep your functions that read/write to files as simple as possible, and place the complex logic in separate functions that can be tested with minimal or zero mocking.

Tagged
Python
pytest