diff options
Diffstat (limited to 'Lib/unittest/mock.py')
-rw-r--r-- | Lib/unittest/mock.py | 18 |
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) |