Skip to content

Commit 72b1004

Browse files
tirkarthicjw296
authored andcommitted
bpo-25597: Ensure wraps' return value is used for magic methods in MagicMock (#16029)
1 parent 10355ed commit 72b1004

File tree

3 files changed

+56
-0
lines changed

3 files changed

+56
-0
lines changed

Lib/unittest/mock.py

+6
Original file line numberDiff line numberDiff line change
@@ -2033,6 +2033,12 @@ def __aiter__():
20332033

20342034

20352035
def _set_return_value(mock, method, name):
2036+
# If _mock_wraps is present then attach it so that wrapped object
2037+
# is used for return value is used when called.
2038+
if mock._mock_wraps is not None:
2039+
method._mock_wraps = getattr(mock._mock_wraps, name)
2040+
return
2041+
20362042
fixed = _return_values.get(name, DEFAULT)
20372043
if fixed is not DEFAULT:
20382044
method.return_value = fixed

Lib/unittest/test/testmock/testmock.py

+47
Original file line numberDiff line numberDiff line change
@@ -715,6 +715,53 @@ def method(self): pass
715715
self.assertRaises(StopIteration, mock.method)
716716

717717

718+
def test_magic_method_wraps_dict(self):
719+
data = {'foo': 'bar'}
720+
721+
wrapped_dict = MagicMock(wraps=data)
722+
self.assertEqual(wrapped_dict.get('foo'), 'bar')
723+
self.assertEqual(wrapped_dict['foo'], 'bar')
724+
self.assertTrue('foo' in wrapped_dict)
725+
726+
# return_value is non-sentinel and takes precedence over wrapped value.
727+
wrapped_dict.get.return_value = 'return_value'
728+
self.assertEqual(wrapped_dict.get('foo'), 'return_value')
729+
730+
# return_value is sentinel and hence wrapped value is returned.
731+
wrapped_dict.get.return_value = sentinel.DEFAULT
732+
self.assertEqual(wrapped_dict.get('foo'), 'bar')
733+
734+
self.assertEqual(wrapped_dict.get('baz'), None)
735+
with self.assertRaises(KeyError):
736+
wrapped_dict['baz']
737+
self.assertFalse('bar' in wrapped_dict)
738+
739+
data['baz'] = 'spam'
740+
self.assertEqual(wrapped_dict.get('baz'), 'spam')
741+
self.assertEqual(wrapped_dict['baz'], 'spam')
742+
self.assertTrue('baz' in wrapped_dict)
743+
744+
del data['baz']
745+
self.assertEqual(wrapped_dict.get('baz'), None)
746+
747+
748+
def test_magic_method_wraps_class(self):
749+
750+
class Foo:
751+
752+
def __getitem__(self, index):
753+
return index
754+
755+
def __custom_method__(self):
756+
return "foo"
757+
758+
759+
klass = MagicMock(wraps=Foo)
760+
obj = klass()
761+
self.assertEqual(obj.__getitem__(2), 2)
762+
self.assertEqual(obj.__custom_method__(), "foo")
763+
764+
718765
def test_exceptional_side_effect(self):
719766
mock = Mock(side_effect=AttributeError)
720767
self.assertRaises(AttributeError, mock)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Ensure, if ``wraps`` is supplied to :class:`unittest.mock.MagicMock`, it is used
2+
to calculate return values for the magic methods instead of using the default
3+
return values. Patch by Karthikeyan Singaravelan.

0 commit comments

Comments
 (0)