Mocking of 'Open' as a Context Manager Made Simple In Python

Using open as a context manager is a great way to ensure your file handles are closed properly and is becoming common:

with open('/some/path', 'w') as f:
    f.write('something')

The issue is that even if you mock out the call to open it is the returned object that is used as a context manager (and has __enter__ and __exit__ called).

Using MagicMock from the mock library, we can mock out context managers very simply. However, mocking open is fiddly enough that a helper function is useful. Here mock_open creates and configures a MagicMock that behaves as a file context manager.

from mock import inPy3k, MagicMock

if inPy3k:
    file_spec = ['_CHUNK_SIZE', '__enter__', '__eq__', '__exit__',
        '__format__', '__ge__', '__gt__', '__hash__', '__iter__', '__le__',
        '__lt__', '__ne__', '__next__', '__repr__', '__str__',
        '_checkClosed', '_checkReadable', '_checkSeekable',
        '_checkWritable', 'buffer', 'close', 'closed', 'detach',
        'encoding', 'errors', 'fileno', 'flush', 'isatty',
        'line_buffering', 'mode', 'name',
        'newlines', 'peek', 'raw', 'read', 'read1', 'readable',
        'readinto', 'readline', 'readlines', 'seek', 'seekable', 'tell',
        'truncate', 'writable', 'write', 'writelines']
else:
    file_spec = file

def mock_open(mock=None, data=None):
    if mock is None:
        mock = MagicMock(spec=file_spec)

    handle = MagicMock(spec=file_spec)
    handle.write.return_value = None
    if data is None:
        handle.__enter__.return_value = handle
    else:
        handle.__enter__.return_value = data
    mock.return_value = handle
    return mock

>>> m = mock_open()
>>> with patch('__main__.open', m, create=True):
...     with open('foo', 'w') as h:
...         h.write('some stuff')
...
>>> m.assert_called_once_with('foo', 'w')
>>> m.mock_calls
[call('foo', 'w'),
 call().__enter__(),
 call().write('some stuff'),
 call().__exit__(None, None, None)]
>>> handle = m()
>>> handle.write.assert_called_once_with('some stuff')

And for reading files, using a StringIO to represent the file handle:

>>> from StringIO import StringIO
>>> m = mock_open(data=StringIO('foo bar baz'))
>>> with patch('__main__.open', m, create=True):
...     with open('foo') as h:
...         result = h.read()
...
>>> m.assert_called_once_with('foo')
>>> assert result == 'foo bar baz'

Note that the StringIO will only be used for the data if open is used as a context manager. If you just configure and use mocks they will work whichever way open is used.

This helper function will be built into mock 0.9.


Source:  http://www.voidspace.org.uk/python/weblog/arch_d7_2012_01_07.shtml