bpo-30541: Add new method to seal mocks (GH61923)

The new method allows the developer to control when to stop the
feature of mocks that automagically creates new mocks when accessing
an attribute that was not declared before

Signed-off-by: Mario Corchero <mariocj89@gmail.com>

2 files changed, 222 insertions, 2 deletions

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index 94e3442..9302ded 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -18,6 +18,7 @@ __all__ = (
     'NonCallableMagicMock',
     'mock_open',
     'PropertyMock',
+    'seal',
 )
 
 
@@ -382,6 +383,7 @@ class NonCallableMock(Base):
         __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 @@ class NonCallableMock(Base):
         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 @@ class NonCallableMock(Base):
             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 @@ class NonCallableMock(Base):
         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 @@ class NonCallableMock(Base):
                 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 @@ class PropertyMock(Mock):
         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
new file mode 100644
index 0000000..0e72b32
--- /dev/null
+++ 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()