Index: Python-3.3.0/Lib/unittest/mock.py =================================================================== --- Python-3.3.0.orig/Lib/unittest/mock.py +++ Python-3.3.0/Lib/unittest/mock.py @@ -2163,6 +2163,24 @@ FunctionTypes = ( 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=''): """ @@ -2173,9 +2191,15 @@ 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 _remaining_data_as_list(): + return list(handle._data) + + def _remaining_data_as_string(): + return ''.join(handle._data) + global file_spec if file_spec is None: import _io @@ -2186,8 +2210,12 @@ def mock_open(mock=None, read_data=''): handle = MagicMock(spec=file_spec) handle.write.return_value = None + handle._data = _iterate_read_data(read_data) handle.__enter__.return_value = handle - handle.read.return_value = read_data + handle.read.side_effect = _remaining_data_as_string + + handle.readline.side_effect = handle._data + handle.readlines.side_effect = _remaining_data_as_list mock.return_value = handle return mock Index: Python-3.3.0/Lib/unittest/test/testmock/testwith.py =================================================================== --- Python-3.3.0.orig/Lib/unittest/test/testmock/testwith.py +++ Python-3.3.0/Lib/unittest/test/testmock/testwith.py @@ -171,6 +171,60 @@ 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_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') if __name__ == '__main__': unittest.main() Index: Python-3.3.0/Doc/library/unittest.mock.rst =================================================================== --- Python-3.3.0.orig/Doc/library/unittest.mock.rst +++ Python-3.3.0/Doc/library/unittest.mock.rst @@ -1950,8 +1950,12 @@ mock_open 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`, `readline`, and `readlines` methods + of the file handle to return. Calls to those methods will take data from + `read_data` until it is depleted. The mock of these methods is pretty + simplistic. If you need more control over the data that you are feeding to + the tested code you will need to customize this mock for yourself. + `read_data` is an empty string by default. Using `open` as a context manager is a great way to ensure your file handles are closed properly and is becoming common::