diff options
author | Michael Foord <michael@voidspace.org.uk> | 2014-04-14 16:09:42 -0400 |
---|---|---|
committer | Michael Foord <michael@voidspace.org.uk> | 2014-04-14 16:09:42 -0400 |
commit | 01bafdcccc4ab653f70379a58a3183fac36e7132 (patch) | |
tree | e82e59083815c2a56bbf276fffe7434fe7daa256 /Lib/unittest/mock.py | |
parent | 061cb3b04d77e490489f64f230fc4dd7a190de4d (diff) | |
download | cpython-git-01bafdcccc4ab653f70379a58a3183fac36e7132.tar.gz |
Issue 17826. Setting an iterable side_effect on a mock created by create_autospec now works
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 a23df5a73d..f1f65221d4 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) |