diff options
author | Michael Foord <michael@voidspace.org.uk> | 2014-04-14 20:10:02 (GMT) |
---|---|---|
committer | Michael Foord <michael@voidspace.org.uk> | 2014-04-14 20:10:02 (GMT) |
commit | fba913f77aa53dff416753c987bf590762c98406 (patch) | |
tree | c9b5cfe56db2e7f87f5b631ce1a5f4e9bf7cd65e /Lib/unittest/mock.py | |
parent | 4f7b0c3c358fe614042951a4ecacf73f08254699 (diff) | |
parent | 01bafdcccc4ab653f70379a58a3183fac36e7132 (diff) | |
download | cpython-fba913f77aa53dff416753c987bf590762c98406.zip cpython-fba913f77aa53dff416753c987bf590762c98406.tar.gz cpython-fba913f77aa53dff416753c987bf590762c98406.tar.bz2 |
Merge
Diffstat (limited to 'Lib/unittest/mock.py')
-rw-r--r-- | Lib/unittest/mock.py | 15 |
1 files changed, 13 insertions, 2 deletions
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index a23df5a..f1f6522 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -343,7 +343,14 @@ def _check_and_set_parent(parent, value, name, new_name): value._mock_name = name return True - +# Internal class to identify if we wrapped an iterator object or not. +class _MockIter(object): + def __init__(self, obj): + self.obj = iter(obj) + def __iter__(self): + return self + def __next__(self): + return next(self.obj) class Base(object): _mock_return_value = DEFAULT @@ -495,7 +502,11 @@ class NonCallableMock(Base): delegated = self._mock_delegate if delegated is None: return self._mock_side_effect - return delegated.side_effect + sf = delegated.side_effect + if sf is not None and not callable(sf) and not isinstance(sf, _MockIter): + sf = _MockIter(sf) + delegated.side_effect = sf + return sf def __set_side_effect(self, value): value = _try_iter(value) |