python · vstinner · Oct 17, 2017 · Jun 1, 2017 · Jun 1, 2017 · Jun 1, 2017
diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst
@@ -2365,3 +2365,23 @@ alternative object as the *autospec* argument:
    a mocked class to create a mock instance *does not* create a real instance.
    It is only attribute lookups - along with calls to :func:`dir` - that are done.
 
+Sealing mocks
+~~~~~~~~~~~~~
+
+.. function:: seal(mock)
+
+    Seal will disable the creation of mock children by preventing to get or set
+    any new attribute on the sealed mock. The sealing process is performed recursively.
+
+    If a mock instance is assigned to an attribute instead of being dynamically created
+    it wont be considered in the sealing chain. This allows to prevent seal from fixing
+    part of the mock object.
+
+        >>> mock = Mock()
+        >>> mock.submock.attribute1 = 2
+        >>> mock.not_submock = mock.Mock()
+        >>> seal(mock)
+        >>> mock.submock.attribute2  # This will raise AttributeError.
+        >>> mock.not_submock.attribute2  # This won't raise.
+
+    .. versionadded:: 3.7
diff --git a/Doc/whatsnew/3.7.rst b/Doc/whatsnew/3.7.rst
@@ -177,6 +177,11 @@ The :const:`~unittest.mock.sentinel` attributes now preserve their identity
 when they are :mod:`copied <copy>` or :mod:`pickled <pickle>`. (Contributed by
 Serhiy Storchaka in :issue:`20804`.)
 
+New function :const:`~unittest.mock.seal` will disable the creation of mock
+children by preventing to get or set any new attribute on the sealed mock.
+The sealing process is performed recursively. (Contributed by Mario Corchero
+in :issue:`30541`.)
+
 xmlrpc.server
 -------------
 

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
@@ -18,6 +18,7 @@
     'NonCallableMagicMock',
     'mock_open',
     'PropertyMock',
+    'seal',
 )
 
 
@@ -382,6 +383,7 @@ def __init__(
         __dict__['_mock_name'] = name
         __dict__['_mock_new_name'] = _new_name
         __dict__['_mock_new_parent'] = _new_parent
+        __dict__['_mock_sealed'] = False
 
         if spec_set is not None:
             spec = spec_set
@@ -608,7 +610,7 @@ def __getattr__(self, name):
         return result
 
 
-    def __repr__(self):
+    def _extract_mock_name(self):
         _name_list = [self._mock_new_name]
         _parent = self._mock_new_parent
         last = self
@@ -638,7 +640,10 @@ def __repr__(self):
             if _name_list[1] not in ('()', '().'):
                 _first += '.'
         _name_list[0] = _first
-        name = ''.join(_name_list)
+        return ''.join(_name_list)
+
+    def __repr__(self):
+        name = self._extract_mock_name()
 
         name_string = ''
         if name not in ('mock', 'mock.'):
@@ -705,6 +710,11 @@ def __setattr__(self, name, value):
         else:
             if _check_and_set_parent(self, value, name, name):
                 self._mock_children[name] = value
+
+        if self._mock_sealed and not hasattr(self, name):
+            mock_name = f'{self._extract_mock_name()}.{name}'
+            raise AttributeError(f'Cannot set {mock_name}')
+
         return object.__setattr__(self, name, value)
 
 
@@ -888,6 +898,12 @@ def _get_child_mock(self, **kw):
                 klass = Mock
         else:
             klass = _type.__mro__[1]
+
+        if self._mock_sealed:
+            attribute = "." + kw["name"] if "name" in kw else "()"
+            mock_name = self._extract_mock_name() + attribute
+            raise AttributeError(mock_name)
+
         return klass(**kw)
 
 
@@ -2401,3 +2417,26 @@ def __get__(self, obj, obj_type):
         return self()
     def __set__(self, obj, val):
         self(val)
+
+
+def seal(mock):
+    """Disable the automatic generation of "submocks"
+
+    Given an input Mock, seals it to ensure no further mocks will be generated
+    when accessing an attribute that was not already defined.
+
+    Submocks are defined as all mocks which were created DIRECTLY from the
+    parent. If a mock is assigned to an attribute of an existing mock,
+    it is not considered a submock.
+
+    """
+    mock._mock_sealed = True
+    for attr in dir(mock):
+        try:
+            m = getattr(mock, attr)
+        except AttributeError:
+            continue
+        if not isinstance(m, NonCallableMock):
+            continue
+        if m._mock_new_parent is mock:
+            seal(m)
diff --git a/Lib/unittest/test/testmock/testsealable.py b/Lib/unittest/test/testmock/testsealable.py
@@ -0,0 +1,181 @@
+import unittest
+from unittest import mock
+
+
+class SampleObject:
+    def __init__(self):
+        self.attr_sample1 = 1
+        self.attr_sample2 = 1
+
+    def method_sample1(self):
+        pass
+
+    def method_sample2(self):
+        pass
+
+
+class TestSealable(unittest.TestCase):
+
+    def test_attributes_return_more_mocks_by_default(self):
+        m = mock.Mock()
+
+        self.assertIsInstance(m.test, mock.Mock)
+        self.assertIsInstance(m.test(), mock.Mock)
+        self.assertIsInstance(m.test().test2(), mock.Mock)
+
+    def test_new_attributes_cannot_be_accessed_on_seal(self):
+        m = mock.Mock()
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError):
+            m.test
+        with self.assertRaises(AttributeError):
+            m()
+
+    def test_new_attributes_cannot_be_set_on_seal(self):
+        m = mock.Mock()
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError):
+            m.test = 1
+
+    def test_existing_attributes_can_be_set_on_seal(self):
+        m = mock.Mock()
+        m.test.test2 = 1
+
+        mock.seal(m)
+        m.test.test2 = 2
+        self.assertEqual(m.test.test2, 2)
+
+    def test_new_attributes_cannot_be_set_on_child_of_seal(self):
+        m = mock.Mock()
+        m.test.test2 = 1
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError):
+            m.test.test3 = 1
+
+    def test_existing_attributes_allowed_after_seal(self):
+        m = mock.Mock()
+
+        m.test.return_value = 3
+
+        mock.seal(m)
+        self.assertEqual(m.test(), 3)
+
+    def test_initialized_attributes_allowed_after_seal(self):
+        m = mock.Mock(test_value=1)
+
+        mock.seal(m)
+        self.assertEqual(m.test_value, 1)
+
+    def test_call_on_sealed_mock_fails(self):
+        m = mock.Mock()
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError):
+            m()
+
+    def test_call_on_defined_sealed_mock_succeeds(self):
+        m = mock.Mock(return_value=5)
+
+        mock.seal(m)
+        self.assertEqual(m(), 5)
+
+    def test_seals_recurse_on_added_attributes(self):
+        m = mock.Mock()
+
+        m.test1.test2().test3 = 4
+
+        mock.seal(m)
+        self.assertEqual(m.test1.test2().test3, 4)
+        with self.assertRaises(AttributeError):
+            m.test1.test2().test4
+        with self.assertRaises(AttributeError):
+            m.test1.test3
+
+    def test_seals_recurse_on_magic_methods(self):
+        m = mock.MagicMock()
+
+        m.test1.test2["a"].test3 = 4
+        m.test1.test3[2:5].test3 = 4
+
+        mock.seal(m)
+        self.assertEqual(m.test1.test2["a"].test3, 4)
+        self.assertEqual(m.test1.test2[2:5].test3, 4)
+        with self.assertRaises(AttributeError):
+            m.test1.test2["a"].test4
+        with self.assertRaises(AttributeError):
+            m.test1.test3[2:5].test4
+
+    def test_seals_dont_recurse_on_manual_attributes(self):
+        m = mock.Mock(name="root_mock")
+
+        m.test1.test2 = mock.Mock(name="not_sealed")
+        m.test1.test2.test3 = 4
+
+        mock.seal(m)
+        self.assertEqual(m.test1.test2.test3, 4)
+        m.test1.test2.test4  # Does not raise
+        m.test1.test2.test4 = 1  # Does not raise
+
+    def test_integration_with_spec_att_definition(self):
+        """You are not restricted when using mock with spec"""
+        m = mock.Mock(SampleObject)
+
+        m.attr_sample1 = 1
+        m.attr_sample3 = 3
+
+        mock.seal(m)
+        self.assertEqual(m.attr_sample1, 1)
+        self.assertEqual(m.attr_sample3, 3)
+        with self.assertRaises(AttributeError):
+            m.attr_sample2
+
+    def test_integration_with_spec_method_definition(self):
+        """You need to defin the methods, even if they are in the spec"""
+        m = mock.Mock(SampleObject)
+
+        m.method_sample1.return_value = 1
+
+        mock.seal(m)
+        self.assertEqual(m.method_sample1(), 1)
+        with self.assertRaises(AttributeError):
+            m.method_sample2()
+
+    def test_integration_with_spec_method_definition_respects_spec(self):
+        """You cannot define methods out of the spec"""
+        m = mock.Mock(SampleObject)
+
+        with self.assertRaises(AttributeError):
+            m.method_sample3.return_value = 3
+
+    def test_sealed_exception_has_attribute_name(self):
+        m = mock.Mock()
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError) as cm:
+            m.SECRETE_name
+        self.assertIn("SECRETE_name", str(cm.exception))
+
+    def test_attribute_chain_is_maintained(self):
+        m = mock.Mock(name="mock_name")
+        m.test1.test2.test3.test4
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError) as cm:
+            m.test1.test2.test3.test4.boom
+        self.assertIn("mock_name.test1.test2.test3.test4.boom", str(cm.exception))
+
+    def test_call_chain_is_maintained(self):
+        m = mock.Mock()
+        m.test1().test2.test3().test4
+
+        mock.seal(m)
+        with self.assertRaises(AttributeError) as cm:
+            m.test1().test2.test3().test4()
+        self.assertIn("mock.test1().test2.test3().test4", str(cm.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2017-10-17-12-04-37.bpo-30541.q3BM6C.rst b/Misc/NEWS.d/next/Library/2017-10-17-12-04-37.bpo-30541.q3BM6C.rst
@@ -0,0 +1,2 @@
+Add new function to seal a mock and prevent the automatically creation of
+child mocks. Patch by Mario Corchero.
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Add new function to seal a mock and prevent the automatically creation of
		child mocks. Patch by Mario Corchero.