summaryrefslogtreecommitdiffstats
path: root/Lib/unittest/mock.py
diff options
context:
space:
mode:
Diffstat (limited to 'Lib/unittest/mock.py')
-rw-r--r--Lib/unittest/mock.py18
1 files changed, 18 insertions, 0 deletions
diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index b49030a..a647e5d 100644
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -1820,6 +1820,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()
@@ -1831,6 +1837,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)