Merge pull request #117 from kaste/better-patch

Author: kaste

CHANGES.txt

@@ -82,6 +82,10 @@ Release 2.0.0
       mock.do(1, 2, x=3)
       assert call.value == ((1, 2), {"x": 3})
 
+- Added `patch_attr` and `patch_dict` for non-callable monkeypatch-style use cases
+  (e.g. `sys.stdout`, `sys.argv`, and environment/config dictionaries) with
+  context-manager support and restoration through `unstub`.
+
 
 
 

docs/index.rst

@@ -78,6 +78,14 @@ State-of-the-art, high-five argument matchers::
     when(math).sqrt(not_(number)).thenRaise(
         TypeError('argument must be a number'))
 
+Captors::
+
+    args, kwargs = captor(), captor()
+    when(mamma).said(*args, **kwargs)
+    # use it ...
+    assert args.value == ("Knock", "You", "Out")
+
+
 No need to `verify` (`assert_called_with`) all the time::
 
     # Different arguments, different answers
@@ -120,6 +128,13 @@ Full async/await support::
     when(module_under_test).http_get('https://example.com', ...).thenReturn('Yep!')
 
 
+Convenience::
+
+    with patch_attr("sys.argv", ["foo", "bar"]):
+    with patch_attr("sys.stdout", StringIO()) as stdout: ...
+    with patch_dict(os.environ, {"user": "bob"}): ...
+
+
 Read
 ----
 

docs/the-functions.rst

@@ -4,11 +4,13 @@
 The functions
 =============
 
-Stable entrypoints are: :func:`when`, :func:`mock`, :func:`unstub`, :func:`verify`, :func:`spy`. New function introduced in v1 are: :func:`when2`, :func:`expect`, :func:`verifyExpectedInteractions`, :func:`verifyStubbedInvocationsAreUsed`, :func:`patch`
+Stable entrypoints are: :func:`when`, :func:`mock`, :func:`unstub`, :func:`verify`, :func:`spy`. New function introduced in v1 are: :func:`when2`, :func:`expect`, :func:`verifyExpectedInteractions`, :func:`verifyStubbedInvocationsAreUsed`, :func:`patch`, :func:`patch_attr`, :func:`patch_dict`
 
 .. autofunction:: when
 .. autofunction:: when2
 .. autofunction:: patch
+.. autofunction:: patch_attr
+.. autofunction:: patch_dict
 .. autofunction:: expect
 .. autofunction:: mock
 .. autofunction:: unstub

mockito/__init__.py

@@ -25,6 +25,8 @@
     when,
     when2,
     patch,
+    patch_attr,
+    patch_dict,
     expect,
     unstub,
     forget_invocations,
@@ -61,6 +63,8 @@
     'when',
     'when2',
     'patch',
+    'patch_attr',
+    'patch_dict',
     'expect',
     'ensureNoUnverifiedInteractions',
     'verify',

mockito/mocking.py

@@ -31,6 +31,7 @@
 from . import invocation, signature, utils
 from . import verification as verificationModule
 from .mock_registry import mock_registry
+from .patching import Patch, patcher
 
 
 __all__ = ['mock']
@@ -41,8 +42,6 @@
 )
 SUPPORTS_MARKCOROUTINEFUNCTION = hasattr(inspect, "markcoroutinefunction")
 
-_MISSING_ATTRIBUTE = object()
-
 _CONFIG_ASYNC_PREFIX = "async "
 _ASYNC_BY_PROTOCOL_METHODS = {"__aenter__", "__aexit__", "__anext__"}
 
@@ -324,7 +323,7 @@ def __init__(
         self.stubbed_invocations: deque[invocation.StubbedInvocation] = deque()
 
         self._original_methods: dict[str, object | None] = {}
-        self._methods_to_unstub: dict[str, object] = {}
+        self._methods_to_unstub: dict[str, Patch] = {}
         self._signatures_store: dict[str, signature.Signature | None] = {}
         self._property_access_context: \
             list[tuple[str, object | None, object]] = []
@@ -447,30 +446,13 @@ def _get_original_method_before_stub(
         if self.spec is None:
             return None, False
 
-        try:
-            return self.spec.__dict__[method_name], True
-        except (AttributeError, KeyError):
-            # If the attr is not directly in __dict__, class specs should use
-            # static lookup so inherited descriptors are preserved as
-            # descriptors (instead of triggering __get__ via getattr).
-            if inspect.isclass(self.spec):
-                try:
-                    return inspect.getattr_static(self.spec, method_name), False
-                except AttributeError:
-                    # If static lookup misses (e.g. metaclass __getattr__),
-                    # fall back to dynamic lookup.
-                    pass
-
-            # For instance specs, keep dynamic getattr so existing
-            # bound-method/spying behavior stays unchanged.
-            return getattr(self.spec, method_name, None), False
-
-    def set_method(self, method_name: str, new_method: object) -> None:
-        setattr(self.mocked_obj, method_name, new_method)
+        return utils.get_original_attribute(self.spec, method_name, default=None)
 
     def replace_method(
-        self, method_name: str, original_method: object | None
-    ) -> None:
+        self,
+        method_name: str,
+        original_method: object | None,
+    ) -> Patch:
         discard_first_arg = self._takes_implicit_self_or_cls(original_method)
 
         def new_mocked_method(*args, **kwargs):
@@ -506,44 +488,36 @@ def new_mocked_method(*args, **kwargs):
         ):
             new_mocked_method = staticmethod(new_mocked_method)
 
-        self.set_method(method_name, new_mocked_method)
+        return patcher.patch_attribute(
+            self.mocked_obj,
+            method_name,
+            new_mocked_method,
+            allow_unstub_by_replacement=False,
+        )
 
     def stub(self, method_name: str) -> None:
         try:
             self._methods_to_unstub[method_name]
         except KeyError:
-            (
-                original_method,
-                was_in_spec
-            ) = self._get_original_method_before_stub(method_name)
-            if was_in_spec:
-                # This indicates the original method was found directly on
-                # the spec object and should therefore be restored by unstub
-                self._methods_to_unstub[method_name] = original_method
-            else:
-                self._methods_to_unstub[method_name] = _MISSING_ATTRIBUTE
-
+            original_method, _ = self._get_original_method_before_stub(method_name)
             self._original_methods[method_name] = original_method
-            self.replace_method(method_name, original_method)
+            self._methods_to_unstub[method_name] = self.replace_method(
+                method_name,
+                original_method,
+            )
 
     def stub_property(self, method_name: str) -> None:
         try:
             self._methods_to_unstub[method_name]
         except KeyError:
-            (
-                original_method,
-                was_in_spec
-            ) = self._get_original_method_before_stub(method_name)
-
+            original_method, _ = self._get_original_method_before_stub(method_name)
             self._original_methods[method_name] = original_method
-            self.set_method(method_name, _mocked_property(self, method_name))
-
-            if was_in_spec:
-                # This indicates the original method was found directly on
-                # the spec object and should therefore be restored by unstub
-                self._methods_to_unstub[method_name] = original_method
-            else:
-                self._methods_to_unstub[method_name] = _MISSING_ATTRIBUTE
+            self._methods_to_unstub[method_name] = patcher.patch_attribute(
+                self.mocked_obj,
+                method_name,
+                _mocked_property(self, method_name),
+                allow_unstub_by_replacement=False,
+            )
 
 
     def forget_stubbed_invocation(
@@ -558,26 +532,18 @@ def forget_stubbed_invocation(
             inv.method_name == invocation.method_name
             for inv in self.stubbed_invocations
         ):
-            original_method = self._methods_to_unstub.pop(
-                invocation.method_name
-            )
-            self.restore_method(invocation.method_name, original_method)
+            patch = self._methods_to_unstub.pop(invocation.method_name)
+            patch.restore_and_unregister()
 
         if self.stubbed_invocations:
             return
 
         mock_registry.unstub(self.mocked_obj)
 
-    def restore_method(self, method_name: str, original_method: object) -> None:
-        if original_method is _MISSING_ATTRIBUTE:
-            delattr(self.mocked_obj, method_name)
-        else:
-            self.set_method(method_name, original_method)
-
     def unstub(self) -> None:
         while self._methods_to_unstub:
-            method_name, original_method = self._methods_to_unstub.popitem()
-            self.restore_method(method_name, original_method)
+            _, patch = self._methods_to_unstub.popitem()
+            patch.restore_and_unregister()
         self.stubbed_invocations = deque()
         self.invocations = []
         self._methods_marked_as_coroutine = set()

mockito/mockito.py

@@ -19,15 +19,16 @@
 # THE SOFTWARE.
 
 from __future__ import annotations
+from collections.abc import Iterable, MutableMapping
 import operator
-from typing import Iterable
 
 from . import invocation
 from . import verification
 
 from .utils import deprecated, get_obj, get_obj_attr_tuple
 from .mocking import Chain, Mock
 from .mock_registry import mock_registry
+from .patching import restore_patch_contextmanager, patcher
 from .verification import VerificationError
 
 
@@ -304,6 +305,100 @@ def patch(fn, attr_or_replacement, replacement=None):
             theMock, name, strict=False)(Ellipsis).thenAnswer(replacement)
 
 
+def patch_attr(obj_or_path, attr_or_replacement, replacement=OMITTED):
+    """Patch/replace an attribute with a concrete value.
+
+    Unlike :func:`patch`, this does *not* record interactions and does not
+    expose verification. It is intended for simple attribute replacement like
+    ``sys.stdout`` or ``sys.argv``.
+
+    Two ways to call this. Either::
+
+        patch_attr('sys.stdout', StringIO())  # two arguments
+        # OR
+        patch_attr(sys, 'stdout', StringIO())  # three arguments
+
+    ``with`` context management is supported and restores the original value
+    on ``__exit__``. ``__enter__`` returns the replacement object.
+
+    .. note:: You must :func:`unstub` after patching, or use `with`
+        statement.
+
+    """
+    if replacement is OMITTED:
+        replacement = attr_or_replacement
+        obj, name = get_obj_attr_tuple(obj_or_path)
+    else:
+        obj, name = obj_or_path, attr_or_replacement
+
+    patch = patcher.patch_attribute(
+        obj,
+        name,
+        replacement,
+        allow_unstub_by_replacement=True,
+    )
+    return restore_patch_contextmanager(patch, replacement)
+
+
+def patch_dict(mapping_or_path, values=None, *, clear=False, remove=None, **kwargs):
+    """Patch/update a dict-like object in place.
+
+    This is a convenience function for test-time dictionary patching,
+    especially for mutable global maps like ``os.environ``.
+
+    Usage::
+
+        patch_dict(os.environ, {'USER': 'foo'})
+        patch_dict(os.environ, [('USER', 'foo')])
+        patch_dict(os.environ, USER='foo')
+        patch_dict(os.environ, remove={'USER', 'PATH'})
+        patch_dict(os.environ, remove=all)
+        patch_dict(os.environ, clear=True)
+        patch_dict('os.environ', {'USER': 'foo'})
+
+    ``with`` context management is supported and restores the original mapping
+    state on ``__exit__``. ``__enter__`` returns the patched mapping.
+
+    ``values`` can be any value accepted by ``dict(values)``.
+    ``kwargs`` are merged into ``values`` and take precedence.
+
+    .. note:: You must :func:`unstub` after patching, or use `with`
+        statement.
+
+    """
+    mapping = (
+        get_obj(mapping_or_path)
+        if isinstance(mapping_or_path, str)
+        else mapping_or_path
+    )
+
+    if not isinstance(mapping, MutableMapping):
+        raise TypeError("target must be a mutable mapping")
+
+    if remove is all:
+        clear = True
+        remove = None
+
+    normalized_remove: tuple[object, ...]
+    if remove is None:
+        normalized_remove = ()
+    elif isinstance(remove, (str, bytes)):
+        normalized_remove = (remove,)
+    elif not isinstance(remove, Iterable):
+        raise TypeError("remove must be iterable, all, or None")
+    else:
+        normalized_remove = tuple(remove)
+
+    updates = {} if values is None else dict(values)
+    updates.update(kwargs)
+    patch = patcher.patch_dictionary(
+        mapping,
+        updates,
+        clear=clear,
+        remove=normalized_remove,
+    )
+    return restore_patch_contextmanager(patch, mapping)
+
 
 def expect(obj, strict=True,
            times=None, atleast=None, atmost=None, between=None):
@@ -348,7 +443,7 @@ def __getattr__(self, method_name):
 
 
 def unstub(*objs):
-    """Unstubs all stubbed methods and functions
+    """Unstubs all stubbed methods, functions, and patched attributes.
 
     If you don't pass in any argument, *all* registered mocks and
     patched modules, classes etc. will be unstubbed.
@@ -363,8 +458,10 @@ def unstub(*objs):
             if isinstance(obj, str):
                 obj = get_obj(obj)
             mock_registry.unstub(obj)
+            patcher.unstub_matching(obj)
     else:
         mock_registry.unstub_all()
+        patcher.unstub_all()
 
 
 def forget_invocations(*objs):