python · claudiubelu · Jun 7, 2017 · Jun 7, 2017
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
@@ -102,6 +102,8 @@ def checksig(_mock_self, *args, **kwargs):
     _copy_func_details(func, checksig)
     type(mock)._mock_check_sig = checksig
 
+    return sig
+
 
 def _copy_func_details(func, funcopy):
     # we explicitly don't copy func.__dict__ into this copy as it would
@@ -152,6 +154,10 @@ def _set_signature(mock, original, instance=False):
     if not _callable(original):
         return
 
+    if '_spec_signature' in mock.__dict__:
+        # mock already has a spec signature. Nothing to do.
+        return mock
+
     skipfirst = isinstance(original, type)
     result = _get_signature_object(original, instance, skipfirst)
     if result is None:
@@ -372,7 +378,8 @@ def __new__(cls, *args, **kw):
     def __init__(
             self, spec=None, wraps=None, name=None, spec_set=None,
             parent=None, _spec_state=None, _new_name='', _new_parent=None,
-            _spec_as_instance=False, _eat_self=None, unsafe=False, **kwargs
+            _spec_as_instance=False, _eat_self=None, unsafe=False,
+            autospec=None, **kwargs
         ):
         if _new_parent is None:
             _new_parent = parent
@@ -382,10 +389,15 @@ def __init__(
         __dict__['_mock_name'] = name
         __dict__['_mock_new_name'] = _new_name
         __dict__['_mock_new_parent'] = _new_parent
+        __dict__['_autospec'] = autospec
 
         if spec_set is not None:
             spec = spec_set
             spec_set = True
+        if autospec is not None:
+            # autospec is even stricter than spec_set.
+            spec = autospec
+            autospec = True
         if _eat_self is None:
             _eat_self = parent is not None
 
@@ -426,12 +438,18 @@ def attach_mock(self, mock, attribute):
         setattr(self, attribute, mock)
 
 
-    def mock_add_spec(self, spec, spec_set=False):
+    def mock_add_spec(self, spec, spec_set=False, autospec=None):
         """Add a spec to a mock. `spec` can either be an object or a
         list of strings. Only attributes on the `spec` can be fetched as
         attributes from the mock.
 
-        If `spec_set` is True then only attributes on the spec can be set."""
+        If `spec_set` is True then only attributes on the spec can be set.
+        If `autospec` is True then only attributes on the spec can be accessed
+        and set, and if a method in the `spec` is called, it's signature is
+        checked.
+        """
+        if autospec is not None:
+            self.__dict__['_autospec'] = autospec
         self._mock_add_spec(spec, spec_set)
 
 
@@ -445,9 +463,9 @@ def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,
                 _spec_class = spec
             else:
                 _spec_class = _get_class(spec)
-            res = _get_signature_object(spec,
-                                        _spec_as_instance, _eat_self)
-            _spec_signature = res and res[1]
+
+            _spec_signature = _check_signature(spec, self, _eat_self,
+                                               _spec_as_instance)
 
             spec = dir(spec)
 
@@ -592,9 +610,20 @@ def __getattr__(self, name):
                 # execution?
                 wraps = getattr(self._mock_wraps, name)
 
+            kwargs = {}
+            if self.__dict__.get('_autospec') is not None:
+                # get the mock's spec attribute with the same name and
+                # pass it to the child.
+                spec_class = self.__dict__.get('_spec_class')
+                spec = getattr(spec_class, name, None)
+                is_type = isinstance(spec_class, type)
+                eat_self = _must_skip(spec_class, name, is_type)
+                kwargs['_eat_self'] = eat_self
+                kwargs['autospec'] = spec
+
             result = self._get_child_mock(
                 parent=self, name=name, wraps=wraps, _new_name=name,
-                _new_parent=self
+                _new_parent=self, **kwargs
             )
             self._mock_children[name]  = result
 
@@ -1313,11 +1342,21 @@ def __enter__(self):
             if original is DEFAULT:
                 raise TypeError("Can't use 'autospec' with create=True")
             spec_set = bool(spec_set)
+            is_class = isinstance(self.target, type)
             if autospec is True:
-                autospec = original
+                if is_class:
+                    # if it's a class, "original" could represent a class
+                    # or static method object (non-callable), instead of the
+                    # method itself, causing a NonCallableMagicMock to be
+                    # created.
+                    autospec = getattr(self.target, self.attribute, original)
+                else:
+                    autospec = original
 
+            eat_self = _must_skip(self.target, self.attribute, is_class)
             new = create_autospec(autospec, spec_set=spec_set,
-                                  _name=self.attribute, **kwargs)
+                                  _name=self.attribute, _eat_self=eat_self,
+                                  **kwargs)
         elif kwargs:
             # can't set keyword args when we aren't creating the mock
             # XXXX If new is a Mock we could call new.configure_mock(**kwargs)
@@ -2211,9 +2250,9 @@ def create_autospec(spec, spec_set=False, instance=False, _parent=None,
             new = _SpecState(original, spec_set, mock, entry, instance)
             mock._mock_children[entry] = new
         else:
+            # we are looping over a spec's attributes, so mock is the parent
+            # of mocked spec attributes.
             parent = mock
-            if isinstance(spec, FunctionTypes):
-                parent = mock.mock
 
             skipfirst = _must_skip(spec, entry, is_type)
             kwargs['_eat_self'] = skipfirst

diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py
@@ -510,6 +510,48 @@ def test_only_allowed_methods_exist(self):
                 )
 
 
+    def _check_autospeced_something(self, something):
+        for method_name in ['meth', 'cmeth', 'smeth']:
+            mock_method = getattr(something, method_name)
+
+            # check that the methods are callable with correct args.
+            mock_method(sentinel.a, sentinel.b, sentinel.c)
+            mock_method(sentinel.a, sentinel.b, sentinel.c, d=sentinel.d)
+            mock_method.assert_has_calls([
+                call(sentinel.a, sentinel.b, sentinel.c),
+                call(sentinel.a, sentinel.b, sentinel.c, d=sentinel.d)])
+
+            # assert that TypeError is raised if the method signature is not
+            # respected.
+            self.assertRaises(TypeError, mock_method)
+            self.assertRaises(TypeError, mock_method, sentinel.a)
+            self.assertRaises(TypeError, mock_method, a=sentinel.a)
+            self.assertRaises(TypeError, mock_method, sentinel.a, sentinel.b,
+                              sentinel.c, e=sentinel.e)
+
+        # assert that AttributeError is raised if the method does not exist.
+        self.assertRaises(AttributeError, getattr, something, 'foolish')
+
+
+    def test_mock_autospec_all_members(self):
+        for spec in [Something, Something()]:
+            mock = Mock(autospec=spec)
+            self._check_autospeced_something(mock)
+
+
+    def test_mock_spec_function(self):
+        def foo(lish):
+            pass
+
+        mock = Mock(spec=foo)
+
+        mock(sentinel.lish)
+        mock.assert_called_once_with(sentinel.lish)
+        self.assertRaises(TypeError, mock)
+        self.assertRaises(TypeError, sentinel.foo, sentinel.lish)
+        self.assertRaises(TypeError, mock, foo=sentinel.foo)
+
+
     def test_from_spec(self):
         class Something(object):
             x = 3
@@ -748,6 +790,21 @@ def test_filter_dir(self):
             patcher.stop()
 
 
+    def test_patch_autospec_obj(self):
+        something = Something()
+        with patch.multiple(something, meth=DEFAULT, cmeth=DEFAULT,
+                            smeth=DEFAULT, autospec=True):
+            self._check_autospeced_something(something)
+
+
+    @patch.object(Something, 'smeth', autospec=True)
+    @patch.object(Something, 'cmeth', autospec=True)
+    @patch.object(Something, 'meth', autospec=True)
+    def test_patch_autospec_class(self, mock_meth, mock_cmeth, mock_smeth):
+        something = Something()
+        self._check_autospeced_something(something)
+
+
     def test_configure_mock(self):
         mock = Mock(foo='bar')
         self.assertEqual(mock.foo, 'bar')
@@ -1376,6 +1433,13 @@ def test_mock_add_spec_magic_methods(self):
             self.assertRaises(TypeError, lambda: mock['foo'])
 
 
+    def test_mock_add_spec_autospec_all_members(self):
+        for spec in [Something, Something()]:
+            mock = Mock()
+            mock.mock_add_spec(spec, autospec=True)
+            self._check_autospeced_something(mock)
+
+
     def test_adding_child_mock(self):
         for Klass in NonCallableMock, Mock, MagicMock, NonCallableMagicMock:
             mock = Klass()