diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst index 8346cb922ee..fcd77277a01 100644 --- a/Doc/library/unittest.mock.rst +++ b/Doc/library/unittest.mock.rst @@ -1995,9 +1995,11 @@ mock_open :meth:`~io.IOBase.readline`, and :meth:`~io.IOBase.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. + simplistic: every time the *mock* is called, the *read_data* is rewound to + the start. 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. When that + is insufficient, one of the in-memory filesystem packages on `PyPI + `_ can offer a realistic filesystem for testing. 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 74f918a5f34..3fbe846fc68 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -2278,6 +2278,24 @@ def mock_open(mock=None, read_data=''): `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(_state[0]) + + def _read_side_effect(*args, **kwargs): + if handle.read.return_value is not None: + return handle.read.return_value + return ''.join(_state[0]) + + def _readline_side_effect(): + if handle.readline.return_value is not None: + while True: + yield handle.readline.return_value + for line in _state[0]: + yield line + + global file_spec if file_spec is None: import _io @@ -2286,42 +2304,31 @@ def mock_open(mock=None, read_data=''): if mock is None: mock = MagicMock(name='open', spec=open) - def make_handle(*args, **kwargs): - # Arg checking is handled by __call__ - def _readlines_side_effect(*args, **kwargs): - if handle.readlines.return_value is not None: - return handle.readlines.return_value - return list(_data) + handle = MagicMock(spec=file_spec) + handle.__enter__.return_value = handle - def _read_side_effect(*args, **kwargs): - if handle.read.return_value is not None: - return handle.read.return_value - return ''.join(_data) + _state = [_iterate_read_data(read_data), None] - 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 + handle.write.return_value = None + handle.read.return_value = None + handle.readline.return_value = None + handle.readlines.return_value = None - handle = MagicMock(spec=file_spec) - handle.__enter__.return_value = handle + handle.read.side_effect = _read_side_effect + _state[1] = _readline_side_effect() + handle.readline.side_effect = _state[1] + handle.readlines.side_effect = _readlines_side_effect - _data = _iterate_read_data(read_data) + def reset_data(*args, **kwargs): + _state[0] = _iterate_read_data(read_data) + if handle.readline.side_effect == _state[1]: + # Only reset the side effect if the user hasn't overridden it. + _state[1] = _readline_side_effect() + handle.readline.side_effect = _state[1] + return DEFAULT - 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 - _check_and_set_parent(mock, handle, None, '()') - return handle - - mock.side_effect = make_handle + mock.side_effect = reset_data + mock.return_value = handle return mock diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py index 32703e6d7bc..976c40fc456 100644 --- a/Lib/unittest/test/testmock/testmock.py +++ b/Lib/unittest/test/testmock/testmock.py @@ -1,5 +1,6 @@ import copy import sys +import tempfile import unittest from unittest.test.testmock.support import is_instance @@ -1329,8 +1330,29 @@ class MockTest(unittest.TestCase): def test_mock_open_reuse_issue_21750(self): mocked_open = mock.mock_open(read_data='data') f1 = mocked_open('a-name') + f1_data = f1.read() f2 = mocked_open('another-name') - self.assertEqual(f1.read(), f2.read()) + f2_data = f2.read() + self.assertEqual(f1_data, f2_data) + + def test_mock_open_write(self): + # Test exception in file writing write() + mock_namedtemp = mock.mock_open(mock.MagicMock(name='JLV')) + with mock.patch('tempfile.NamedTemporaryFile', mock_namedtemp): + mock_filehandle = mock_namedtemp.return_value + mock_write = mock_filehandle.write + mock_write.side_effect = OSError('Test 2 Error') + def attempt(): + tempfile.NamedTemporaryFile().write('asd') + self.assertRaises(OSError, attempt) + + def test_mock_open_alter_readline(self): + mopen = mock.mock_open(read_data='foo\nbarn') + mopen.return_value.readline.side_effect = lambda *args:'abc' + first = mopen().readline() + second = mopen().readline() + self.assertEqual('abc', first) + self.assertEqual('abc', second) def test_mock_parents(self): for Klass in Mock, MagicMock: diff --git a/Lib/unittest/test/testmock/testwith.py b/Lib/unittest/test/testmock/testwith.py index ddcfe77e4d9..b6bfb754ad6 100644 --- a/Lib/unittest/test/testmock/testwith.py +++ b/Lib/unittest/test/testmock/testwith.py @@ -141,6 +141,7 @@ class TestMockOpen(unittest.TestCase): def test_mock_open_context_manager(self): mock = mock_open() + handle = mock.return_value with patch('%s.open' % __name__, mock, create=True): with open('foo') as f: f.read() @@ -148,8 +149,7 @@ class TestMockOpen(unittest.TestCase): expected_calls = [call('foo'), call().__enter__(), call().read(), call().__exit__(None, None, None)] self.assertEqual(mock.mock_calls, expected_calls) - # mock_open.return_value is no longer static, because - # readline support requires that it mutate state + self.assertIs(f, handle) def test_mock_open_context_manager_multiple_times(self): mock = mock_open()