Issue 17826. Setting an iterable side_effect on a mock created by create_autospec now works
diff --git a/NEWS b/NEWS
index 8fe2cf4..af528f3 100644
--- a/NEWS
+++ b/NEWS
@@ -1,3 +1,6 @@
+- Issue #17826: setting an iterable side_effect on a mock function created by
+ create_autospec now works. Patch by Kushal Das.
+
- Issue #20968: unittest.mock.MagicMock now supports division.
Patch by Johannes Baiter.
diff --git a/mock.py b/mock.py
index 6a9c901..d47c3f1 100644
--- a/mock.py
+++ b/mock.py
@@ -61,6 +61,9 @@
from six import wraps
+inPy3k = sys.version_info[0] == 3
+
+
try:
inspectsignature = inspect.signature
except AttributeError:
@@ -102,11 +105,15 @@
# Python 2.4 compatibility
BaseException = Exception
-try:
- next
-except NameError:
- def next(obj):
- return obj.next()
+if not inPy3k:
+ # Python 2's next() can't handle a non-iterator with a __next__ method.
+ _next = next
+ def next(obj, _next=_next):
+ if getattr(obj, '__next__', None):
+ return obj.__next__()
+ return _next(obj)
+
+ del _next
BaseExceptions = (BaseException,)
@@ -127,9 +134,6 @@
return False
return regex.match(string)
-
-inPy3k = sys.version_info[0] == 3
-
self = 'im_self'
builtin = '__builtin__'
if inPy3k:
@@ -465,7 +469,14 @@
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
@@ -617,7 +628,11 @@
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/tests/testmock.py b/tests/testmock.py
index d03da18..fd184db 100644
--- a/tests/testmock.py
+++ b/tests/testmock.py
@@ -178,6 +178,24 @@
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')
