From 01bafdcccc4ab653f70379a58a3183fac36e7132 Mon Sep 17 00:00:00 2001
From: Michael Foord <michael@voidspace.org.uk>
Date: Mon, 14 Apr 2014 16:09:42 -0400
Subject: Issue 17826. Setting an iterable side_effect on a mock created by
 create_autospec now works

---
 Lib/unittest/mock.py                   | 15 +++++++++++++--
 Lib/unittest/test/testmock/testmock.py | 18 ++++++++++++++++++
 Misc/NEWS                              |  3 +++
 3 files changed, 34 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)
diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py
index 20cc654..70f19d3 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 2dedeb4..09c0f40 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -33,6 +33,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.
 
-- 
cgit v0.12