bpo-25597: Ensure wraps' return value is used for magic methods in MagicMock (#16029)

2 files changed, 53 insertions, 0 deletions

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index beed717..1acafc5 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -2033,6 +2033,12 @@ _side_effect_methods = {
 
 
 def _set_return_value(mock, method, name):
+    # If _mock_wraps is present then attach it so that wrapped object
+    # is used for return value is used when called.
+    if mock._mock_wraps is not None:
+        method._mock_wraps = getattr(mock._mock_wraps, name)
+        return
+
     fixed = _return_values.get(name, DEFAULT)
     if fixed is not DEFAULT:
         method.return_value = fixed
diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py
index 1329346..6773467 100644
--- a/Lib/unittest/test/testmock/testmock.py
+++ b/Lib/unittest/test/testmock/testmock.py
@@ -715,6 +715,53 @@ class MockTest(unittest.TestCase):
         self.assertRaises(StopIteration, mock.method)
 
 
+    def test_magic_method_wraps_dict(self):
+        data = {'foo': 'bar'}
+
+        wrapped_dict = MagicMock(wraps=data)
+        self.assertEqual(wrapped_dict.get('foo'), 'bar')
+        self.assertEqual(wrapped_dict['foo'], 'bar')
+        self.assertTrue('foo' in wrapped_dict)
+
+        # return_value is non-sentinel and takes precedence over wrapped value.
+        wrapped_dict.get.return_value = 'return_value'
+        self.assertEqual(wrapped_dict.get('foo'), 'return_value')
+
+        # return_value is sentinel and hence wrapped value is returned.
+        wrapped_dict.get.return_value = sentinel.DEFAULT
+        self.assertEqual(wrapped_dict.get('foo'), 'bar')
+
+        self.assertEqual(wrapped_dict.get('baz'), None)
+        with self.assertRaises(KeyError):
+            wrapped_dict['baz']
+        self.assertFalse('bar' in wrapped_dict)
+
+        data['baz'] = 'spam'
+        self.assertEqual(wrapped_dict.get('baz'), 'spam')
+        self.assertEqual(wrapped_dict['baz'], 'spam')
+        self.assertTrue('baz' in wrapped_dict)
+
+        del data['baz']
+        self.assertEqual(wrapped_dict.get('baz'), None)
+
+
+    def test_magic_method_wraps_class(self):
+
+        class Foo:
+
+            def __getitem__(self, index):
+                return index
+
+            def __custom_method__(self):
+                return "foo"
+
+
+        klass = MagicMock(wraps=Foo)
+        obj = klass()
+        self.assertEqual(obj.__getitem__(2), 2)
+        self.assertEqual(obj.__custom_method__(), "foo")
+
+
     def test_exceptional_side_effect(self):
         mock = Mock(side_effect=AttributeError)
         self.assertRaises(AttributeError, mock)