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>

5 files changed, 249 insertions, 2 deletions

diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst
index 9e8bf11..6fdfdc4 100644
--- 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
index a0c20a0..11b99e2 100644
--- a/Doc/whatsnew/3.7.rst
+++ b/Doc/whatsnew/3.7.rst
@@ -304,6 +304,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
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()
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
new file mode 100644
index 0000000..7eb5e16
--- /dev/null
+++ 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.