8000 bpo-31855: unittest.mock.mock_open() results now respects the argumen… · lisroach/cpython@11a8832 · GitHub
[go: up one dir, main page]

Skip to content

Commit 11a8832

Browse files
Rémi Lapeyrecjw296
Rémi Lapeyre
authored andcommitted
bpo-31855: unittest.mock.mock_open() results now respects the argument of read([size]) (pythonGH-11521)
unittest.mock.mock_open() results now respects the argument of read([size]) Co-Authored-By: remilapeyre <remi.lapeyre@henki.fr>
1 parent ad4ed87 commit 11a8832

Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
__version__ = '1.0'
2626

2727

28+
import io
2829
import inspect
2930
import pprint
3031
import sys
@@ -2318,25 +2319,12 @@ def __init__(self, spec, spec_set=False, parent=None,
23182319

23192320
file_spec = None
23202321

2321-
def _iterate_read_data(read_data):
2322-
# Helper for mock_open:
2323-
# Retrieve lines from read_data via a generator so that separate calls to
2324-
# readline, read, and readlines are properly interleaved
2325-
sep = b'\n' if isinstance(read_data, bytes) else '\n'
2326-
data_as_list = [l + sep for l in read_data.split(sep)]
2327-
2328-
if data_as_list[-1] == sep:
2329 10000 -
# If the last line ended in a newline, the list comprehension will have an
2330-
# extra entry that's just a newline. Remove this.
2331-
data_as_list = data_as_list[:-1]
2332-
else:
2333-
# If there wasn't an extra newline by itself, then the file being
2334-
# emulated doesn't have a newline to end the last line remove the
2335-
# newline that our naive format() added
2336-
data_as_list[-1] = data_as_list[-1][:-1]
23372322

2338-
for line in data_as_list:
2339-
yield line
2323+
def _to_stream(read_data):
2324+
if isinstance(read_data, bytes):
2325+
return io.BytesIO(read_data)
2326+
else:
2327+
return io.StringIO(read_data)
23402328

23412329

23422330
def mock_open(mock=None, read_data=''):
@@ -2351,20 +2339,23 @@ def mock_open(mock=None, read_data=''):
23512339
`read_data` is a string for the `read`, `readline` and `readlines` of the
23522340
file handle to return. This is an empty string by default.
23532341
"""
2342+
_read_data = _to_stream(read_data)
2343+
_state = [_read_data, None]
2344+
23542345
def _readlines_side_effect(*args, **kwargs):
23552346
if handle.readlines.return_value is not None:
23562347
return handle.readlines.return_value
2357-
return list(_state[0])
2348+
return _state[0].readlines(*args, **kwargs)
23582349

23592350
def _read_side_effect(*args, **kwargs):
23602351
if handle.read.return_value is not None:
23612352
return handle.read.return_value
2362-
return type(read_data)().join(_state[0])
2353+
return _state[0].read(*args, **kwargs)
23632354

2364-
def _readline_side_effect():
2355+
def _readline_side_effect(*args, **kwargs):
23652356
yield from _iter_side_effect()
23662357
while True:
2367-
yield type(read_data)()
2358+
yield _state[0].readline(*args, **kwargs)
23682359

23692360
def _iter_side_effect():
23702361
if handle.readline.return_value is not None:
@@ -2384,8 +2375,6 @@ def _iter_side_effect():
23842375
handle = MagicMock(spec=file_spec)
23852376
handle.__enter__.return_value = handle
23862377

2387-
_state = [_iterate_read_data(read_data), None]
2388-
23892378
handle.write.return_value = None
23902379
handle.read.return_value = None
23912380
handle.readline.return_value = None
@@ -2398,7 +2387,7 @@ def _iter_side_effect():
23982387
handle.__iter__.side_effect = _iter_side_effect
23992388

24002389
def reset_data(*args, **kwargs):
2401-
_state[0] = _iterate_read_data(read_data)
2390+
_state[0] = _to_stream(read_data)
24022391
if handle.readline.side_effect == _state[1]:
24032392
# Only reset the side effect if the user hasn't overridden it.
24042393
_state[1] = _readline_side_effect()
Original file line numberDiff line numberDiff line change
@@ -283,7 +283,12 @@ def test_mock_open_read_with_argument(self):
283283
# for mocks returned by mock_open
284284
some_data = 'foo\nbar\nbaz'
285285
mock = mock_open(read_data=some_data)
286-
self.assertEqual(mock().read(10), some_data)
286+
self.assertEqual(mock().read(10), some_data[:10])
287+
self.assertEqual(mock().read(10), some_data[:10])
288+
289+
f = mock()
290+
self.assertEqual(f.read(10), some_data[:10])
291+
self.assertEqual(f.read(10), some_data[10:])
287292

288293

289294
def test_interleaved_reads(self):
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
:func:`unittest.mock.mock_open` results now respects the argument of read([size]).
2+
Patch contributed by Rémi Lapeyre.