Skip to content

Commit

Permalink
bpo-31855: unittest.mock.mock_open() results now respects the argumen…
Browse files Browse the repository at this point in the history
…t 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>
(cherry picked from commit 11a8832)

Co-authored-by: Rémi Lapeyre <remi.lapeyre@henki.fr>
  • Loading branch information
Rémi Lapeyre authored and miss-islington committed May 7, 2019
1 parent ffa29b5 commit e461be5
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 26 deletions.
39 changes: 14 additions & 25 deletions Lib/unittest/mock.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
__version__ = '1.0'


import io
import inspect
import pprint
import sys
Expand Down Expand Up @@ -2331,25 +2332,12 @@ def __init__(self, spec, spec_set=False, parent=None,

file_spec = None

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
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] == 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]
else:
# If there wasn't an extra newline by itself, then the file being
# emulated doesn't have a newline to end the last line remove the
# newline that our naive format() added
data_as_list[-1] = data_as_list[-1][:-1]

for line in data_as_list:
yield line
def _to_stream(read_data):
if isinstance(read_data, bytes):
return io.BytesIO(read_data)
else:
return io.StringIO(read_data)


def mock_open(mock=None, read_data=''):
Expand All @@ -2364,20 +2352,23 @@ def mock_open(mock=None, read_data=''):
`read_data` is a string for the `read`, `readline` and `readlines` of the
file handle to return. This is an empty string by default.
"""
_read_data = _to_stream(read_data)
_state = [_read_data, None]

def _readlines_side_effect(*args, **kwargs):
if handle.readlines.return_value is not None:
return handle.readlines.return_value
return list(_state[0])
return _state[0].readlines(*args, **kwargs)

def _read_side_effect(*args, **kwargs):
if handle.read.return_value is not None:
return handle.read.return_value
return type(read_data)().join(_state[0])
return _state[0].read(*args, **kwargs)

def _readline_side_effect():
def _readline_side_effect(*args, **kwargs):
yield from _iter_side_effect()
while True:
yield type(read_data)()
yield _state[0].readline(*args, **kwargs)

def _iter_side_effect():
if handle.readline.return_value is not None:
Expand All @@ -2397,8 +2388,6 @@ def _iter_side_effect():
handle = MagicMock(spec=file_spec)
handle.__enter__.return_value = handle

_state = [_iterate_read_data(read_data), None]

handle.write.return_value = None
handle.read.return_value = None
handle.readline.return_value = None
Expand All @@ -2411,7 +2400,7 @@ def _iter_side_effect():
handle.__iter__.side_effect = _iter_side_effect

def reset_data(*args, **kwargs):
_state[0] = _iterate_read_data(read_data)
_state[0] = _to_stream(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()
Expand Down
7 changes: 6 additions & 1 deletion Lib/unittest/test/testmock/testwith.py
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,12 @@ def test_mock_open_read_with_argument(self):
# for mocks returned by mock_open
some_data = 'foo\nbar\nbaz'
mock = mock_open(read_data=some_data)
self.assertEqual(mock().read(10), some_data)
self.assertEqual(mock().read(10), some_data[:10])
self.assertEqual(mock().read(10), some_data[:10])

f = mock()
self.assertEqual(f.read(10), some_data[:10])
self.assertEqual(f.read(10), some_data[10:])


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

0 comments on commit e461be5

Please sign in to comment.