gh-74772: Adds autospec argument for Mock / MagicMock / AsyncMock

Lib/test/test_unittest/testmock/testmock.py

@@ -1,9 +1,10 @@
+import asyncio
 import copy
 import re
 import sys
 import tempfile
 
-from test.support import ALWAYS_EQ
+from test.support import ALWAYS_EQ, requires_working_socket
 import unittest
 from test.test_unittest.testmock.support import is_instance
 from unittest import mock
@@ -38,6 +39,16 @@ def cmeth(cls, a, b, c, d=None): pass
     def smeth(a, b, c, d=None): pass
 
 
+class SomethingAsync(object):
+    async def meth(self, a, b, c, d=None): pass
+
+    @classmethod
+    async def cmeth(cls, a, b, c, d=None): pass
+
+    @staticmethod
+    async def smeth(a, b, c, d=None): pass
+
+
 class SomethingElse(object):
     def __init__(self):
         self._instance = None
@@ -695,6 +706,72 @@ def test_attributes(mock):
         test_attributes(Mock(spec=Something()))
 
 
+    def _check_autospeced_something(self, something):
+        # assert that AttributeError is raised if the method does not exist.
+        self.assertRaises(AttributeError, getattr, something, 'foolish')
+
+        self._check_autospeced_something_method(something.meth)
+        self._check_autospeced_something_method(something.cmeth)
+        self._check_autospeced_something_method(something.smeth)
+
+
+    def _check_autospeced_something_method(self, mock_method):
+        # 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)
+
+
+    def test_mock_autospec_all_members(self):
+        for spec in [Something, Something()]:
+            mock_something = Mock(autospec=spec)
+            self._check_autospeced_something(mock_something)
+
+
+    def _check_autospeced_something_async(self, something):
+        # assert that AttributeError is raised if the method does not exist.
+        self.assertRaises(AttributeError, getattr, something, 'foolish')
+
+        self._check_autospeced_something_method_async(something.meth)
+        self._check_autospeced_something_method_async(something.cmeth)
+        self._check_autospeced_something_method_async(something.smeth)
+
+
+    def _check_autospeced_something_method_async(self, mock_method):
+        # check that the methods are callable with correct args.
+        asyncio.run(mock_method(sentinel.a, sentinel.b, sentinel.c))
+        asyncio.run(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)
+
+
+    @requires_working_socket()
+    def test_mock_autospec_all_members_async(self):
+        for spec in [SomethingAsync, SomethingAsync()]:
+            mock_something = AsyncMock(autospec=spec)
+            self._check_autospeced_something_async(mock_something)
+
+
     def test_wraps_calls(self):
         real = Mock()
 
@@ -1974,6 +2051,21 @@ 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_something = Mock()
+            mock_something.mock_add_spec(spec, autospec=True)
+            self._check_autospeced_something(mock_something)
+
+
+    @requires_working_socket()
+    def test_mock_add_spec_autospec_all_members_async(self):
+        for spec in [SomethingAsync, SomethingAsync()]:
+            mock_something = Mock()
+            mock_something.mock_add_spec(spec, autospec=True)
+            self._check_autospeced_something_async(mock_something)
+
+
     def test_adding_child_mock(self):
         for Klass in (NonCallableMock, Mock, MagicMock, NonCallableMagicMock,
                       AsyncMock):

Lib/unittest/mock.py

@@ -136,6 +136,8 @@ def checksig(self, /, *args, **kwargs):
     type(mock)._mock_check_sig = checksig
     type(mock).__signature__ = sig
 
+    return func, sig
+
 
 def _copy_func_details(func, funcopy):
     # we explicitly don't copy func.__dict__ into this copy as it would
@@ -463,7 +465,8 @@ def __new__(
     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
@@ -474,10 +477,15 @@ def __init__(
         __dict__['_mock_new_name'] = _new_name
         __dict__['_mock_new_parent'] = _new_parent
         __dict__['_mock_sealed'] = False
+        __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
 
@@ -520,12 +528,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)
 
 
@@ -543,8 +557,13 @@ def _mock_add_spec(self, spec, spec_set, _spec_as_instance=False,
                 _spec_class = spec
             else:
                 _spec_class = type(spec)
-            res = _get_signature_object(spec,
-                                        _spec_as_instance, _eat_self)
+
+            if self.__dict__.get('_autospec') is None:
+                res = _get_signature_object(spec, _spec_as_instance, _eat_self)
+            else:
+                res = _check_signature(spec, self, _eat_self,
+                                       _spec_as_instance)
+
             _spec_signature = res and res[1]
 
             spec_list = dir(spec)
@@ -712,9 +731,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
 

Misc/NEWS.d/next/Library/2017-11-20-06-36-00.bpo-30587.1SsAy9.rst

@@ -0,0 +1,3 @@
+mock: Added autospec argument to the constructors Mock constructors and mock_add_spec.
+Passing the autospec argument will also check the method signatures of the mocked
+methods.