diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index 9d9b3e2..ad8ca6a 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -2265,7 +2265,18 @@ 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 it's a bytes object, temporarily decode it, because format() is only + # available on str + if type(read_data) == bytes: + read_data = read_data.decode() + + data_as_list = ['{}\n'.format(l).encode() for l in + read_data.split('\n')] + else: + 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 @@ -2300,7 +2311,7 @@ def mock_open(mock=None, read_data=''): def _read_side_effect(*args, **kwargs): if handle.read.return_value is not None: return handle.read.return_value - return ''.join(_data) + return type(read_data)().join(_data) def _readline_side_effect(): if handle.readline.return_value is not None: diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py index 3a104cb..b3e0f90 100644 --- a/Lib/unittest/test/testmock/testmock.py +++ b/Lib/unittest/test/testmock/testmock.py @@ -8,7 +8,7 @@ from unittest.mock import ( call, DEFAULT, patch, sentinel, MagicMock, Mock, NonCallableMock, NonCallableMagicMock, _CallList, - create_autospec + create_autospec, mock_open ) @@ -1444,5 +1444,18 @@ class MockTest(unittest.TestCase): mock.foo + def test_mock_open_read_bytes(self): + m = mock_open(read_data= b'abc') + with patch('%s.open' % __name__, m, create=True) : + with open('abc', 'rb') as f: + self.assertEqual(f.read(), b'abc') + + + def test_mock_open_readlines_bytes(self): + m = mock_open(read_data= b'abc\ndef\nghi') + with patch('%s.open' % __name__, m, create=True) : + with open('abc', 'rb') as f: + self.assertEqual(f.readlines(), [b'abc\n', b'def\n', b'ghi']) + if __name__ == '__main__': unittest.main() diff --git a/Lib/unittest/test/testmock/testwith.py b/Lib/unittest/test/testmock/testwith.py index f54e051..dd97968 100644 --- a/Lib/unittest/test/testmock/testwith.py +++ b/Lib/unittest/test/testmock/testwith.py @@ -139,6 +139,8 @@ class TestMockOpen(unittest.TestCase): mock.assert_called_once_with('foo') + + def test_mock_open_context_manager(self): mock = mock_open() handle = mock.return_value @@ -210,6 +212,20 @@ class TestMockOpen(unittest.TestCase): self.assertEqual(result, ['foo\n', 'bar\n', 'baz']) + def test_read_bytes(self): + m = mock_open(read_data= b'abc') + with patch('%s.open' % __name__, m, create=True) : + with open('abc', 'rb') as f: + self.assertEqual(f.read(), b'abc') + + + def testreadlines_bytes(self): + m = mock_open(read_data= b'abc\ndef\nghi') + with patch('%s.open' % __name__, m, create=True) : + with open('abc', 'rb') as f: + self.assertEqual(f.readlines(), [b'abc\n', b'def\n', b'ghi']) + + def test_mock_open_read_with_argument(self): # At one point calling read with an argument was broken # for mocks returned by mock_open @@ -255,5 +271,6 @@ class TestMockOpen(unittest.TestCase): self.assertEqual(handle.readline(), 'bar') + if __name__ == '__main__': unittest.main()