diff options
Diffstat (limited to 'django/test/utils.py')
| -rw-r--r-- | django/test/utils.py | 23 |
1 files changed, 17 insertions, 6 deletions
diff --git a/django/test/utils.py b/django/test/utils.py index e626667b09..d1f7d19546 100644 --- a/django/test/utils.py +++ b/django/test/utils.py @@ -1,3 +1,4 @@ +import asyncio import logging import re import sys @@ -362,12 +363,22 @@ class TestContextDecorator: raise TypeError('Can only decorate subclasses of unittest.TestCase') def decorate_callable(self, func): - @wraps(func) - def inner(*args, **kwargs): - with self as context: - if self.kwarg_name: - kwargs[self.kwarg_name] = context - return func(*args, **kwargs) + if asyncio.iscoroutinefunction(func): + # If the inner function is an async function, we must execute async + # as well so that the `with` statement executes at the right time. + @wraps(func) + async def inner(*args, **kwargs): + with self as context: + if self.kwarg_name: + kwargs[self.kwarg_name] = context + return await func(*args, **kwargs) + else: + @wraps(func) + def inner(*args, **kwargs): + with self as context: + if self.kwarg_name: + kwargs[self.kwarg_name] = context + return func(*args, **kwargs) return inner def __call__(self, decorated): |