Customizing unittest.mock.mock_open for iteration

Question

How should I customize unittest.mock.mock_open to handle this code?

file: impexpdemo.py
def import_register(register_fn):
    with open(register_fn) as f:
           


        

Answer 1

As of Python 3.6, the mocked file-like object returned by the unittest.mock_open method doesn't support iteration. This bug was reported in 2014 and it is still open as of 2017.

Thus code like this silently yields zero iterations:

f_open = unittest.mock.mock_open(read_data='foo\nbar\n')
f = f_open('blah')
for line in f:
  print(line)

You can work around this limitation via adding a method to the mocked object that returns a proper line iterator:

def mock_open(*args, **kargs):
  f_open = unittest.mock.mock_open(*args, **kargs)
  f_open.return_value.__iter__ = lambda self : iter(self.readline, '')
  return f_open

Answer 2

The mock_open() object does indeed not implement iteration.

If you are not using the file object as a context manager, you could use:

m = unittest.mock.MagicMock(name='open', spec=open)
m.return_value = iter(self.TEST_TEXT)

with unittest.mock.patch('builtins.open', m):

Now open() returns an iterator, something that can be directly iterated over just like a file object can be, and it'll also work with next(). It can not, however, be used as a context manager.

You can combine this with mock_open() then provide a __iter__ and __next__ method on the return value, with the added benefit that mock_open() also adds the prerequisites for use as a context manager:

# Note: read_data must be a string!
m = unittest.mock.mock_open(read_data=''.join(self.TEST_TEXT))
m.return_value.__iter__ = lambda self: self
m.return_value.__next__ = lambda self: next(iter(self.readline, ''))

The return value here is a MagicMock object specced from the file object (Python 2) or the in-memory file objects (Python 3), but only the read, write and __enter__ methods have been stubbed out.

The above doesn't work in Python 2 because a) Python 2 expects next to exist, not __next__ and b) next is not treated as a special method in Mock (rightly so), so even if you renamed __next__ to next in the above example the type of the return value won't have a next method. For most cases it would be enough to make the file object produced an iterable rather than an iterator with:

# Python 2!
m = mock.mock_open(read_data=''.join(self.TEST_TEXT))
m.return_value.__iter__ = lambda self: iter(self.readline, '')

Any code that uses iter(fileobj) will then work (including a for loop).

There is a open issue in the Python tracker that aims to remedy this gap.

Answer 3

I found the following solution:

text_file_data = '\n'.join(["a line here", "the second line", "another 
line in the file"])
with patch('__builtin__.open', mock_open(read_data=text_file_data), 
create=True) as m:
    # mock_open doesn't properly handle iterating over the open file with for line in file:
    # but if we set the return value like this, it works.
    m.return_value.__iter__.return_value = text_file_data.splitlines()
    with open('filename', 'rU') as f:
        for line in f:
            print line