gh-98086: Now ``patch.dict`` can decorate async functions (#98095)

1 files changed, 18 insertions, 0 deletions

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index 096b1a5..a273753 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -1809,6 +1809,12 @@ class _patch_dict(object):
     def __call__(self, f):
         if isinstance(f, type):
             return self.decorate_class(f)
+        if inspect.iscoroutinefunction(f):
+            return self.decorate_async_callable(f)
+        return self.decorate_callable(f)
+
+
+    def decorate_callable(self, f):
         @wraps(f)
         def _inner(*args, **kw):
             self._patch_dict()
@@ -1820,6 +1826,18 @@ class _patch_dict(object):
         return _inner
 
 
+    def decorate_async_callable(self, f):
+        @wraps(f)
+        async def _inner(*args, **kw):
+            self._patch_dict()
+            try:
+                return await f(*args, **kw)
+            finally:
+                self._unpatch_dict()
+
+        return _inner
+
+
     def decorate_class(self, klass):
         for attr in dir(klass):
             attr_value = getattr(klass, attr)