diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index a23df5a73dd..f1f65221d40 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) diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py index 20cc6541e66..70f19d35732 100644 --- a/Lib/unittest/test/testmock/testmock.py +++ b/Lib/unittest/test/testmock/testmock.py @@ -154,6 +154,24 @@ class MockTest(unittest.TestCase): mock = Mock(side_effect=side_effect, return_value=sentinel.RETURN) self.assertEqual(mock(), sentinel.RETURN) + def test_autospec_side_effect(self): + # Test for issue17826 + results = [1, 2, 3] + def effect(): + return results.pop() + def f(): + pass + + mock = create_autospec(f) + mock.side_effect = [1, 2, 3] + self.assertEqual([mock(), mock(), mock()], [1, 2, 3], + "side effect not used correctly in create_autospec") + # Test where side effect is a callable + results = [1, 2, 3] + mock = create_autospec(f) + mock.side_effect = effect + self.assertEqual([mock(), mock(), mock()], [3, 2, 1], + "callable side effect not used correctly") @unittest.skipUnless('java' in sys.platform, 'This test only applies to Jython') diff --git a/Misc/NEWS b/Misc/NEWS index 9c4f013108d..654a7a230e1 100644 --- a/Misc/NEWS +++ b/Misc/NEWS @@ -42,6 +42,9 @@ Core and Builtins Library ------- +- Issue #17826: setting an iterable side_effect on a mock function created by + create_autospec now works. Patch by Kushal Das. + - Issue #7776: Fix ``Host:`` header and reconnection when using http.client.HTTPConnection.set_tunnel(). Patch by Nikolaus Rath.