diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst index ac9dd3b6f793a2..a980f338a443c3 100644 --- a/Doc/library/unittest.mock.rst +++ b/Doc/library/unittest.mock.rst @@ -2095,6 +2095,10 @@ mock_open .. versionchanged:: 3.5 *read_data* is now reset on each call to the *mock*. + .. versionchanged:: 3.8 + Added :meth:`__iter__` to implementation so that iteration (such as in for + loops) correctly consumes *read_data*. + Using :func:`open` as a context manager is a great way to ensure your file handles are closed properly and is becoming common:: diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index 9302dedae7fdad..070fed60e768d1 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -2356,14 +2356,16 @@ def _read_side_effect(*args, **kwargs): return type(read_data)().join(_state[0]) def _readline_side_effect(): + yield from _iter_side_effect() + while True: + yield type(read_data)() + + def _iter_side_effect(): if handle.readline.return_value is not None: while True: yield handle.readline.return_value for line in _state[0]: yield line - while True: - yield type(read_data)() - global file_spec if file_spec is None: @@ -2387,6 +2389,7 @@ def _readline_side_effect(): _state[1] = _readline_side_effect() handle.readline.side_effect = _state[1] handle.readlines.side_effect = _readlines_side_effect + handle.__iter__.side_effect = _iter_side_effect def reset_data(*args, **kwargs): _state[0] = _iterate_read_data(read_data) diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py index b64c8663d21226..c7bfa277b511ca 100644 --- a/Lib/unittest/test/testmock/testmock.py +++ b/Lib/unittest/test/testmock/testmock.py @@ -1450,6 +1450,16 @@ def test_mock_open_reuse_issue_21750(self): f2_data = f2.read() self.assertEqual(f1_data, f2_data) + def test_mock_open_dunder_iter_issue(self): + # Test dunder_iter method generates the expected result and + # consumes the iterator. + mocked_open = mock.mock_open(read_data='Remarkable\nNorwegian Blue') + f1 = mocked_open('a-name') + lines = [line for line in f1] + self.assertEqual(lines[0], 'Remarkable\n') + self.assertEqual(lines[1], 'Norwegian Blue') + self.assertEqual(list(f1), []) + def test_mock_open_write(self): # Test exception in file writing write() mock_namedtemp = mock.mock_open(mock.MagicMock(name='JLV')) diff --git a/Lib/unittest/test/testmock/testwith.py b/Lib/unittest/test/testmock/testwith.py index a7bee730030190..43b36a11995261 100644 --- a/Lib/unittest/test/testmock/testwith.py +++ b/Lib/unittest/test/testmock/testwith.py @@ -188,6 +188,7 @@ def test_read_data(self): def test_readline_data(self): # Check that readline will return all the lines from the fake file + # And that once fully consumed, readline will return an empty string. mock = mock_open(read_data='foo\nbar\nbaz\n') with patch('%s.open' % __name__, mock, create=True): h = open('bar') @@ -197,6 +198,7 @@ def test_readline_data(self): self.assertEqual(line1, 'foo\n') self.assertEqual(line2, 'bar\n') self.assertEqual(line3, 'baz\n') + self.assertEqual(h.readline(), '') # Check that we properly emulate a file that doesn't end in a newline mock = mock_open(read_data='foo') @@ -204,6 +206,19 @@ def test_readline_data(self): h = open('bar') result = h.readline() self.assertEqual(result, 'foo') + self.assertEqual(h.readline(), '') + + + def test_dunder_iter_data(self): + # Check that dunder_iter 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') + lines = [l for l in h] + self.assertEqual(lines[0], 'foo\n') + self.assertEqual(lines[1], 'bar\n') + self.assertEqual(lines[2], 'baz\n') + self.assertEqual(h.readline(), '') def test_readlines_data(self): diff --git a/Misc/NEWS.d/next/Library/2018-04-30-22-43-31.bpo-32933.M3iI_y.rst b/Misc/NEWS.d/next/Library/2018-04-30-22-43-31.bpo-32933.M3iI_y.rst new file mode 100644 index 00000000000000..4de7a8f927d58b --- /dev/null +++ b/Misc/NEWS.d/next/Library/2018-04-30-22-43-31.bpo-32933.M3iI_y.rst @@ -0,0 +1,2 @@ +:func:`unittest.mock.mock_open` now supports iteration over the file +contents. Patch by Tony Flury.