Merge pull request #118 from kaste/partial-unstub

Author: kaste

CHANGES.txt

@@ -86,6 +86,24 @@ Release 2.0.0
   (e.g. `sys.stdout`, `sys.argv`, and environment/config dictionaries) with
   context-manager support and restoration through `unstub`.
 
+  E.g.::
+
+      patch_attr("sys.argv", ["foo", "bar"])
+      with patch_attr("sys.stdout", StringIO()) as stdout: ...
+      with patch_dict(os.environ, {"user": "bob"}): ...
+
+- Added explicit partial-`unstub` targeting by host + attribute name.
+  This complements method-reference partial unstub (e.g. `unstub(cat.meow)`)
+  and supports tuple form and shorthand form, including multiple attributes
+  in one call.
+
+  E.g.::
+
+      unstub(cat.meow)
+      unstub((cat, "meow"))
+      unstub(cat, "meow")
+      unstub((cat, "meow"), (os.path, "exists"))
+
 
 
 

docs/walk-through.rst

@@ -52,7 +52,7 @@ When patching, you **MUST** **not** forget to :func:`unstub` of course! You can
 ::
 
     from mockito import unstub
-    unstub()  # restore os.path module
+    unstub()  # restore all patched/stubbed objects
 
 Usually you do this unconditionally in your `teardown` function. If you're using `pytest`, you could define a fixture instead
 

mockito/mock_registry.py

@@ -88,13 +88,14 @@ def mock_for(self, obj: object) -> Mock | None:
     def obj_for(self, mock: Mock) -> object | None:
         return self.mocks.lookup(mock)
 
-    def unstub(self, obj: object) -> None:
+    def unstub(self, obj: object) -> bool:
         try:
             mock = self.mocks.pop(obj)
         except KeyError:
-            pass
+            return False
         else:
             mock.unstub()
+            return True
 
     def unstub_mock(self, mock: Mock) -> None:
         self.mocks.pop_value(mock)

mockito/mocking.py

@@ -540,6 +540,24 @@ def forget_stubbed_invocation(
 
         mock_registry.unstub(self.mocked_obj)
 
+    def unstub_method(self, method_name: str) -> None:
+        invocations = [
+            invoc
+            for invoc in self.stubbed_invocations
+            if invoc.method_name == method_name
+        ]
+        if not invocations:
+            return
+
+        for invoc in invocations:
+            invoc.forget_self()
+
+        self.invocations = [
+            invocation
+            for invocation in self.invocations
+            if invocation.method_name != method_name
+        ]
+
     def unstub(self) -> None:
         while self._methods_to_unstub:
             _, patch = self._methods_to_unstub.popitem()

mockito/mockito.py

@@ -448,20 +448,142 @@ def unstub(*objs):
     If you don't pass in any argument, *all* registered mocks and
     patched modules, classes etc. will be unstubbed.
 
+    You can also unstub a single method/function target, e.g.::
+
+        unstub(os.path.exists)
+        unstub("os.path.exists")
+        unstub(cat.meow)
+
+    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)
-    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):
+    if not callable(target):
+        return None
+
+    host = getattr(target, "__self__", None)
+    attr_name = getattr(target, "__name__", None)
+    if host is not None and attr_name is not None:
+        return host, attr_name
+
+    target_function = _unwrap_unstub_target(target)
+    for theMock in mock_registry.get_registered_mocks():
+        for method_name, patch in theMock._methods_to_unstub.items():
+            replacement = getattr(patch, "replacement", None)
+            if _unwrap_unstub_target(replacement) is target_function:
+                return theMock.mocked_obj, method_name
+
+    return None
+
+
+
+def _unwrap_unstub_target(target):
+    if isinstance(target, (staticmethod, classmethod)):
+        return target.__func__
+
+    return getattr(target, "__func__", target)
 
 
 def forget_invocations(*objs):

mockito/patching.py

@@ -68,14 +68,28 @@ def patch_dictionary(
         self._register_patch(dict_patch)
         return dict_patch
 
-    def unstub_matching(self, obj: object) -> None:
+    def unstub_matching(self, obj: object) -> bool:
         matching = [
             patch for patch in self._patches
             if patch.matches_unstub_target(obj)
         ]
         for patch in reversed(matching):
             patch.restore_and_unregister()
 
+        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

@@ -45,6 +45,42 @@ def testUnstubsByDottedPath(self):
 
         self.assertEqual(False, os.path.exists("test"))
 
+    def testCanUnstubSingleFunctionByFunctionTarget(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 testCanUnstubSingleFunctionByDottedFunctionPath(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 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/staticmethods_test.py

@@ -50,6 +50,15 @@ def testUnstubs(self):
         unstub()
         self.assertEqual("woof", Dog.bark())
 
+    def testCanUnstubSingleStaticmethodByFunctionTarget(self):
+        when(Dog).bark().thenReturn("miau")
+        when(Dog).barkHardly(1, 2).thenReturn("arf")
+
+        unstub(Dog.bark)
+
+        self.assertEqual("woof", Dog.bark())
+        self.assertEqual("arf", Dog.barkHardly(1, 2))
+
     # TODO decent test case please :) without testing irrelevant implementation
     # details
     def testUnstubShouldPreserveMethodType(self):