1 files changed, 64 insertions, 0 deletions
diff --git a/tests/_autospec.py b/tests/_autospec.py
new file mode 100644
index 000000000..ee2fc1973
--- /dev/null
+++ b/tests/_autospec.py
@@ -0,0 +1,64 @@
+import contextlib
+import functools
+import unittest.mock
+from typing import Callable
+
+
+[email protected](unittest.mock._patch.decoration_helper)
+[email protected]
+def _decoration_helper(self, patched, args, keywargs):
+    """Skips adding patchings as args if their `dont_pass` attribute is True."""
+    # Don't ask what this does. It's just a copy from stdlib, but with the dont_pass check added.
+    extra_args = []
+    with contextlib.ExitStack() as exit_stack:
+        for patching in patched.patchings:
+            arg = exit_stack.enter_context(patching)
+            if not getattr(patching, "dont_pass", False):
+                # Only add the patching as an arg if dont_pass is False.
+                if patching.attribute_name is not None:
+                    keywargs.update(arg)
+                elif patching.new is unittest.mock.DEFAULT:
+                    extra_args.append(arg)
+
+        args += tuple(extra_args)
+        yield args, keywargs
+
+
+[email protected](unittest.mock._patch.copy)
+def _copy(self):
+    """Copy the `dont_pass` attribute along with the standard copy operation."""
+    patcher_copy = _copy.original(self)
+    patcher_copy.dont_pass = getattr(self, "dont_pass", False)
+    return patcher_copy
+
+
+# Monkey-patch the patcher class :)
+_copy.original = unittest.mock._patch.copy
+unittest.mock._patch.copy = _copy
+unittest.mock._patch.decoration_helper = _decoration_helper
+
+
+def autospec(target, *attributes: str, pass_mocks: bool = True, **patch_kwargs) -> Callable:
+    """
+    Patch multiple `attributes` of a `target` with autospecced mocks and `spec_set` as True.
+
+    If `pass_mocks` is True, pass the autospecced mocks as arguments to the decorated object.
+    """
+    # Caller's kwargs should take priority and overwrite the defaults.
+    kwargs = dict(spec_set=True, autospec=True)
+    kwargs.update(patch_kwargs)
+
+    # Import the target if it's a string.
+    # This is to support both object and string targets like patch.multiple.
+    if type(target) is str:
+        target = unittest.mock._importer(target)
+
+    def decorator(func):
+        for attribute in attributes:
+            patcher = unittest.mock.patch.object(target, attribute, **kwargs)
+            if not pass_mocks:
+                # A custom attribute to keep track of which patchings should be skipped.
+                patcher.dont_pass = True
+            func = patcher(func)
+        return func
+    return decorator