diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index 654462cbf7e..b14bf01b28f 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -2680,6 +2680,11 @@ def mock_open(mock=None, read_data=''): for line in _state[0]: yield line + def _next_side_effect(): + if handle.readline.return_value is not None: + return handle.readline.return_value + return next(_state[0]) + global file_spec if file_spec is None: import _io @@ -2701,6 +2706,7 @@ def mock_open(mock=None, read_data=''): handle.readline.side_effect = _state[1] handle.readlines.side_effect = _readlines_side_effect handle.__iter__.side_effect = _iter_side_effect + handle.__next__.side_effect = _next_side_effect def reset_data(*args, **kwargs): _state[0] = _to_stream(read_data) diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py index 307b8b7657a..0f30bccc9cf 100644 --- a/Lib/unittest/test/testmock/testmock.py +++ b/Lib/unittest/test/testmock/testmock.py @@ -1702,6 +1702,19 @@ class MockTest(unittest.TestCase): self.assertEqual(lines[1], 'Norwegian Blue') self.assertEqual(list(f1), []) + def test_mock_open_using_next(self): + mocked_open = mock.mock_open(read_data='1st line\n2nd line\n3rd line') + f1 = mocked_open('a-name') + line1 = next(f1) + line2 = f1.__next__() + lines = [line for line in f1] + self.assertEqual(line1, '1st line\n') + self.assertEqual(line2, '2nd line\n') + self.assertEqual(lines[0], '3rd line') + self.assertEqual(list(f1), []) + with self.assertRaises(StopIteration): + next(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 5172c222d97..42ebf3898c8 100644 --- a/Lib/unittest/test/testmock/testwith.py +++ b/Lib/unittest/test/testmock/testwith.py @@ -230,7 +230,22 @@ class TestMockOpen(unittest.TestCase): self.assertEqual(lines[1], 'bar\n') self.assertEqual(lines[2], 'baz\n') self.assertEqual(h.readline(), '') + with self.assertRaises(StopIteration): + next(h) + def test_next_data(self): + # Check that next will correctly return the next available + # line and plays well with the dunder_iter part. + mock = mock_open(read_data='foo\nbar\nbaz\n') + with patch('%s.open' % __name__, mock, create=True): + h = open('bar') + line1 = next(h) + line2 = next(h) + lines = [l for l in h] + self.assertEqual(line1, 'foo\n') + self.assertEqual(line2, 'bar\n') + self.assertEqual(lines[0], 'baz\n') + self.assertEqual(h.readline(), '') def test_readlines_data(self): # Test that emulating a file that ends in a newline character works diff --git a/Misc/NEWS.d/next/Library/2019-05-22-15-26-08.bpo-37008.WPbv31.rst b/Misc/NEWS.d/next/Library/2019-05-22-15-26-08.bpo-37008.WPbv31.rst new file mode 100644 index 00000000000..42747aead49 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2019-05-22-15-26-08.bpo-37008.WPbv31.rst @@ -0,0 +1,2 @@ +Add support for calling :func:`next` with the mock resulting from +:func:`unittest.mock.mock_open`