diff options
Diffstat (limited to 'Lib')
| -rw-r--r-- | Lib/unittest/mock.py | 10 | ||||
| -rw-r--r-- | Lib/unittest/test/testmock/testmock.py | 10 |
2 files changed, 17 insertions, 3 deletions
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py index 8cedcb4..8378b0f 100644 --- a/Lib/unittest/mock.py +++ b/Lib/unittest/mock.py @@ -891,7 +891,10 @@ class CallableMixin(Base): raise effect if not _callable(effect): - return next(effect) + result = next(effect) + if _is_exception(result): + raise result + return result ret_val = effect(*args, **kwargs) if ret_val is DEFAULT: @@ -931,8 +934,9 @@ class Mock(CallableMixin, NonCallableMock): arguments as the mock, and unless it returns `DEFAULT`, the return value of this function is used as the return value. - Alternatively `side_effect` can be an exception class or instance. In - this case the exception will be raised when the mock is called. + If `side_effect` is an iterable then each call to the mock will return + the next value from the iterable. If any of the members of the iterable + are exceptions they will be raised instead of returned. If `side_effect` is an iterable then each call to the mock will return the next value from the iterable. diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py index ae9822e..64fd1a1 100644 --- a/Lib/unittest/test/testmock/testmock.py +++ b/Lib/unittest/test/testmock/testmock.py @@ -868,6 +868,16 @@ class MockTest(unittest.TestCase): self.assertRaises(StopIteration, mock) + def test_side_effect_iterator_exceptions(self): + for Klass in Mock, MagicMock: + iterable = (ValueError, 3, KeyError, 6) + m = Klass(side_effect=iterable) + self.assertRaises(ValueError, m) + self.assertEqual(m(), 3) + self.assertRaises(KeyError, m) + self.assertEqual(m(), 6) + + def test_side_effect_setting_iterator(self): mock = Mock() mock.side_effect = iter([1, 2, 3]) |