Add explicit unstub target syntax for object attributes

Author: kaste

mockito/mockito.py

@@ -454,38 +454,109 @@ def unstub(*objs):
         unstub("os.path.exists")
         unstub(cat.meow)
 
-    In these cases only that one attribute is restored, while other stubs on
-    the same object stay active.
+    Or explicitly target one attribute by host and name, e.g.::
+
+        unstub((cat, "meow"))
+        unstub(cat, "meow")
+        unstub((cat, "meow"), (os.path, "exists"))
+
+    In these cases only the selected attributes are restored, while other stubs
+    on the same objects stay active.
 
     Note that additionally, the underlying registry will be cleaned.
     After an `unstub` you can't :func:`verify` anymore because all
     interactions will be forgotten.
     """
 
-    if objs:
-        for obj in objs:
-            if isinstance(obj, str):
-                obj = get_obj(obj)
-
-            # mock_registry.unstub(obj)
-            # patcher.unstub_matching(obj)
-            if (
-                mock_registry.unstub(obj)
-                or patcher.unstub_matching(obj)
-            ):
-                return
-
-            resolved_target = _resolve_unstub_attr_target(obj)
-            if resolved_target is None:
-                continue
-
-            host, attr_name = resolved_target
-            host_mock = mock_registry.mock_for(host)
-            if host_mock is not None:
-                host_mock.unstub_method(attr_name)
-    else:
+    if not objs:
         mock_registry.unstub_all()
         patcher.unstub_all()
+        return
+
+    explicit_attr_targets, generic_targets = _partition_unstub_targets(objs)
+
+    for host, attr_name in explicit_attr_targets:
+        _unstub_attr_target(host, attr_name)
+
+    for obj in generic_targets:
+        if isinstance(obj, str):
+            obj = get_obj(obj)
+
+        if mock_registry.unstub(obj) or patcher.unstub_matching(obj):
+            continue
+
+        resolved_target = _resolve_unstub_attr_target(obj)
+        if resolved_target is None:
+            continue
+
+        host, attr_name = resolved_target
+        _unstub_attr_target(host, attr_name)
+
+
+
+def _partition_unstub_targets(objs):
+    if _is_unstub_attr_pair_arguments(objs):
+        host, attr_name = objs
+        return [
+            _normalize_unstub_attr_target(host, attr_name)
+        ], []
+
+    explicit_attr_targets = []
+    generic_targets = []
+
+    for obj in objs:
+        explicit_attr_target = _coerce_unstub_attr_target_tuple(obj)
+        if explicit_attr_target is None:
+            generic_targets.append(obj)
+            continue
+
+        explicit_attr_targets.append(explicit_attr_target)
+
+    return explicit_attr_targets, generic_targets
+
+
+
+def _is_unstub_attr_pair_arguments(objs):
+    return (
+        len(objs) == 2
+        and not isinstance(objs[0], tuple)
+        and _looks_like_attr_name(objs[1])
+    )
+
+
+
+def _coerce_unstub_attr_target_tuple(target):
+    if not isinstance(target, tuple) or len(target) != 2:
+        return None
+
+    host, attr_name = target
+    if not _looks_like_attr_name(attr_name):
+        return None
+
+    return _normalize_unstub_attr_target(host, attr_name)
+
+
+
+def _normalize_unstub_attr_target(host, attr_name):
+    if isinstance(host, str):
+        host = get_obj(host)
+
+    return host, attr_name
+
+
+
+def _looks_like_attr_name(value):
+    return isinstance(value, str) and bool(value) and "." not in value
+
+
+
+def _unstub_attr_target(host, attr_name):
+    host_mock = mock_registry.mock_for(host)
+    if host_mock is not None:
+        host_mock.unstub_method(attr_name)
+
+    patcher.unstub_attribute(host, attr_name)
+
 
 
 def _resolve_unstub_attr_target(target):

mockito/patching.py

@@ -78,6 +78,18 @@ def unstub_matching(self, obj: object) -> bool:
 
         return bool(matching)
 
+    def unstub_attribute(self, obj: object, attr_name: str) -> bool:
+        matching = [
+            patch for patch in self._patches
+            if isinstance(patch, _AttrPatch)
+            if patch.obj is obj
+            if patch.attr_name == attr_name
+        ]
+        for patch in reversed(matching):
+            patch.restore_and_unregister()
+
+        return bool(matching)
+
     def unstub_all(self) -> None:
         for patch in reversed(self._patches.copy()):
             patch.restore_and_unregister()

tests/modulefunctions_test.py

@@ -63,6 +63,24 @@ def testCanUnstubSingleFunctionByDottedFunctionPath(self):
         self.assertEqual(False, os.path.exists("test"))
         self.assertEqual("mocked", os.path.dirname("/tmp/file.txt"))
 
+    def testCanUnstubSingleFunctionByExplicitTargetTuple(self):
+        when(os.path).exists("test").thenReturn(True)
+        when(os.path).dirname(any(str)).thenReturn("mocked")
+
+        unstub((os.path, "exists"))
+
+        self.assertEqual(False, os.path.exists("test"))
+        self.assertEqual("mocked", os.path.dirname("/tmp/file.txt"))
+
+    def testCanUnstubSingleFunctionByExplicitTargetArguments(self):
+        when(os.path).exists("test").thenReturn(True)
+        when(os.path).dirname(any(str)).thenReturn("mocked")
+
+        unstub(os.path, "exists")
+
+        self.assertEqual(False, os.path.exists("test"))
+        self.assertEqual("mocked", os.path.dirname("/tmp/file.txt"))
+
     def testStubs(self):
         when(os.path).exists("test").thenReturn(True)
 

tests/patch_attr_test.py

@@ -56,6 +56,31 @@ def test_patch_attr_can_be_unstubbed_by_replacement_object():
     assert holder.value == "original"
 
 
+def test_patch_attr_can_be_unstubbed_by_explicit_object_and_attribute():
+    holder = Holder()
+
+    patch_attr(holder, "value", "patched")
+    assert holder.value == "patched"
+
+    unstub(holder, "value")
+    assert holder.value == "original"
+
+
+def test_explicit_unstub_attribute_restores_patch_attr_layer_over_when_stub():
+    class LocalHolder:
+        def value(self):
+            return "original"
+
+    holder = LocalHolder()
+
+    when(holder).value().thenReturn("stubbed")
+    patch_attr(holder, "value", lambda: "patched")
+    assert holder.value() == "patched"
+
+    unstub(holder, "value")
+    assert holder.value() == "original"
+
+
 def test_nested_patch_attr_restores_correctly():
     holder = Holder()
 

tests/unstub_test.py

@@ -77,6 +77,51 @@ def testPartialUnstubByMethodReferenceKeepsDetachedChainAlive(self):
 
         assert grand.sleep() == "ok"
 
+    def testPartialUnstubByExplicitTargetTuple(self):
+        cat = mock(strict=True)
+
+        when(cat).meow().thenReturn("Miau")
+        when(cat).runs().thenReturn("Yip")
+
+        unstub((cat, "meow"))
+
+        with pytest.raises(AttributeError):
+            cat.meow()
+
+        assert cat.runs() == "Yip"
+
+    def testPartialUnstubByExplicitTargetArguments(self):
+        cat = mock(strict=True)
+
+        when(cat).meow().thenReturn("Miau")
+        when(cat).runs().thenReturn("Yip")
+
+        unstub(cat, "meow")
+
+        with pytest.raises(AttributeError):
+            cat.meow()
+
+        assert cat.runs() == "Yip"
+
+    def testPartialUnstubByMultipleExplicitTargetTuples(self):
+        cat = mock(strict=True)
+        dog = mock(strict=True)
+
+        when(cat).meow().thenReturn("Miau")
+        when(cat).runs().thenReturn("Yip")
+        when(dog).waggle().thenReturn("Yup")
+        when(dog).bark().thenReturn("Wuff")
+
+        unstub((cat, "meow"), (dog, "waggle"))
+
+        with pytest.raises(AttributeError):
+            cat.meow()
+        with pytest.raises(AttributeError):
+            dog.waggle()
+
+        assert cat.runs() == "Yip"
+        assert dog.bark() == "Wuff"
+
     def testPartialUnstubByMethodReferenceForgetsMethodInvocations(self):
         cat = mock()