diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/unittest/mock.py | 7 | ||||
| -rw-r--r-- | Lib/unittest/test/testmock/testwith.py | 28 |
2 files changed, 32 insertions, 3 deletions
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index efe5763..4c58639 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -2271,9 +2271,10 @@ 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')] + sep = b'\n' if isinstance(read_data, bytes) else '\n' + data_as_list = [l + sep for l in read_data.split(sep)] - if data_as_list[-1] == '\n': + if data_as_list[-1] == sep: # 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] @@ -2307,7 +2308,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(_state[0]) + return type(read_data)().join(_state[0]) def _readline_side_effect(): if handle.readline.return_value is not None: diff --git a/Lib/unittest/test/testmock/testwith.py b/Lib/unittest/test/testmock/testwith.py index b6bfb75..a7bee73 100644 --- a/Lib/unittest/test/testmock/testwith.py +++ b/Lib/unittest/test/testmock/testwith.py @@ -224,6 +224,34 @@ class TestMockOpen(unittest.TestCase): self.assertEqual(result, ['foo\n', 'bar\n', 'baz']) + def test_read_bytes(self): + mock = mock_open(read_data=b'\xc6') + with patch('%s.open' % __name__, mock, create=True): + with open('abc', 'rb') as f: + result = f.read() + self.assertEqual(result, b'\xc6') + + + def test_readline_bytes(self): + m = mock_open(read_data=b'abc\ndef\nghi\n') + with patch('%s.open' % __name__, m, create=True): + with open('abc', 'rb') as f: + line1 = f.readline() + line2 = f.readline() + line3 = f.readline() + self.assertEqual(line1, b'abc\n') + self.assertEqual(line2, b'def\n') + self.assertEqual(line3, b'ghi\n') + + + def test_readlines_bytes(self): + m = mock_open(read_data=b'abc\ndef\nghi\n') + with patch('%s.open' % __name__, m, create=True): + with open('abc', 'rb') as f: + result = f.readlines() + self.assertEqual(result, [b'abc\n', b'def\n', b'ghi\n']) + + def test_mock_open_read_with_argument(self): # At one point calling read with an argument was broken # for mocks returned by mock_open |