Closes issue 17467. Add readline and readlines support to unittest.mock.mock_open

2 files changed, 132 insertions, 6 deletions

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index ea79ae3..2b1effb 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -934,8 +934,6 @@ class CallableMixin(Base):
                 return result
 
             ret_val = effect(*args, **kwargs)
-            if ret_val is DEFAULT:
-                ret_val = self.return_value
 
         if (self._mock_wraps is not None and
              self._mock_return_value is DEFAULT):
@@ -2207,6 +2205,24 @@ MethodWrapperTypes = (
 
 file_spec = None
 
+def _iterate_read_data(read_data):
+    # Helper for mock_open:
+    # Retrieve lines from read_data via a generator so that separate calls to
+    # readline, read, and readlines are properly interleaved
+    data_as_list = ['{}\n'.format(l) for l in read_data.split('\n')]
+
+    if data_as_list[-1] == '\n':
+        # If the last line ended in a newline, the list comprehension will have an
+        # extra entry that's just a newline.  Remove this.
+        data_as_list = data_as_list[:-1]
+    else:
+        # If there wasn't an extra newline by itself, then the file being
+        # emulated doesn't have a newline to end the last line  remove the
+        # newline that our naive format() added
+        data_as_list[-1] = data_as_list[-1][:-1]
+
+    for line in data_as_list:
+        yield line
 
 def mock_open(mock=None, read_data=''):
     """
@@ -2217,9 +2233,27 @@ def mock_open(mock=None, read_data=''):
     default) then a `MagicMock` will be created for you, with the API limited
     to methods or attributes available on standard file handles.
 
-    `read_data` is a string for the `read` method of the file handle to return.
-    This is an empty string by default.
+    `read_data` is a string for the `read` methoddline`, and `readlines` of the
+    file handle to return.  This is an empty string by default.
     """
+    def _readlines_side_effect(*args, **kwargs):
+        if handle.readlines.return_value is not None:
+            return handle.readlines.return_value
+        return list(_data)
+
+    def _read_side_effect(*args, **kwargs):
+        if handle.read.return_value is not None:
+            return handle.read.return_value
+        return ''.join(_data)
+
+    def _readline_side_effect():
+        if handle.readline.return_value is not None:
+            while True:
+                yield handle.readline.return_value
+        for line in _data:
+            yield line
+
+
     global file_spec
     if file_spec is None:
         import _io
@@ -2229,9 +2263,18 @@ def mock_open(mock=None, read_data=''):
         mock = MagicMock(name='open', spec=open)
 
     handle = MagicMock(spec=file_spec)
-    handle.write.return_value = None
     handle.__enter__.return_value = handle
-    handle.read.return_value = read_data
+
+    _data = _iterate_read_data(read_data)
+
+    handle.write.return_value = None
+    handle.read.return_value = None
+    handle.readline.return_value = None
+    handle.readlines.return_value = None
+
+    handle.read.side_effect = _read_side_effect
+    handle.readline.side_effect = _readline_side_effect()
+    handle.readlines.side_effect = _readlines_side_effect
 
     mock.return_value = handle
     return mock
diff --git a/Lib/unittest/test/testmock/testwith.py b/Lib/unittest/test/testmock/testwith.py
index 0a0cfad..f54e051 100644
--- a/Lib/unittest/test/testmock/testwith.py
+++ b/Lib/unittest/test/testmock/testwith.py
@@ -172,5 +172,88 @@ class TestMockOpen(unittest.TestCase):
         self.assertEqual(result, 'foo')
 
 
+    def test_readline_data(self):
+        # Check that readline will return all the lines from the fake file
+        mock = mock_open(read_data='foo\nbar\nbaz\n')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            line1 = h.readline()
+            line2 = h.readline()
+            line3 = h.readline()
+        self.assertEqual(line1, 'foo\n')
+        self.assertEqual(line2, 'bar\n')
+        self.assertEqual(line3, 'baz\n')
+
+        # Check that we properly emulate a file that doesn't end in a newline
+        mock = mock_open(read_data='foo')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            result = h.readline()
+        self.assertEqual(result, 'foo')
+
+
+    def test_readlines_data(self):
+        # Test that emulating a file that ends in a newline character works
+        mock = mock_open(read_data='foo\nbar\nbaz\n')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            result = h.readlines()
+        self.assertEqual(result, ['foo\n', 'bar\n', 'baz\n'])
+
+        # Test that files without a final newline will also be correctly
+        # emulated
+        mock = mock_open(read_data='foo\nbar\nbaz')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            result = h.readlines()
+
+        self.assertEqual(result, ['foo\n', 'bar\n', 'baz'])
+
+
+    def test_mock_open_read_with_argument(self):
+        # At one point calling read with an argument was broken
+        # for mocks returned by mock_open
+        some_data = 'foo\nbar\nbaz'
+        mock = mock_open(read_data=some_data)
+        self.assertEqual(mock().read(10), some_data)
+
+
+    def test_interleaved_reads(self):
+        # Test that calling read, readline, and readlines pulls data
+        # sequentially from the data we preload with
+        mock = mock_open(read_data='foo\nbar\nbaz\n')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            line1 = h.readline()
+            rest = h.readlines()
+        self.assertEqual(line1, 'foo\n')
+        self.assertEqual(rest, ['bar\n', 'baz\n'])
+
+        mock = mock_open(read_data='foo\nbar\nbaz\n')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            line1 = h.readline()
+            rest = h.read()
+        self.assertEqual(line1, 'foo\n')
+        self.assertEqual(rest, 'bar\nbaz\n')
+
+
+    def test_overriding_return_values(self):
+        mock = mock_open(read_data='foo')
+        handle = mock()
+
+        handle.read.return_value = 'bar'
+        handle.readline.return_value = 'bar'
+        handle.readlines.return_value = ['bar']
+
+        self.assertEqual(handle.read(), 'bar')
+        self.assertEqual(handle.readline(), 'bar')
+        self.assertEqual(handle.readlines(), ['bar'])
+
+        # call repeatedly to check that a StopIteration is not propagated
+        self.assertEqual(handle.readline(), 'bar')
+        self.assertEqual(handle.readline(), 'bar')
+
+
 if __name__ == '__main__':
     unittest.main()