diff --git a/mocket/async_mocket.py b/mocket/async_mocket.py index ea3d5330..5628cfa2 100644 --- a/mocket/async_mocket.py +++ b/mocket/async_mocket.py @@ -1,4 +1,4 @@ -from mocket import Mocketizer +from mocket import Mocketizer, mocketize def get_async_mocketize(): @@ -13,7 +13,7 @@ async def __aexit__(*args, **kwargs): def async_wrap(*args, **kwargs): return Mocketizer.wrap(*args, **kwargs) - return AsyncMocketizer.async_wrap + return mocketize(cls=AsyncMocketizer) async_mocketize = get_async_mocketize() diff --git a/mocket/mocket.py b/mocket/mocket.py index a400e77c..f03fcaea 100644 --- a/mocket/mocket.py +++ b/mocket/mocket.py @@ -606,28 +606,27 @@ def check_and_call(self, method): if callable(method): method() - @classmethod - def wrap(cls, test=None, truesocket_recording_dir=None): - def wrapper(t, *args, **kw): - instance = args[0] if args else None - namespace = None - if truesocket_recording_dir: - namespace = ".".join( - ( - instance.__class__.__module__, - instance.__class__.__name__, - t.__name__, - ) - ) - with cls( - instance, - namespace=namespace, - truesocket_recording_dir=truesocket_recording_dir, - ): - t(*args, **kw) - return wrapper - return decorator.decorator(wrapper, test) +def wrapper(test, cls=Mocketizer, truesocket_recording_dir=None, *args, **kw): + instance = args[0] if args else None + namespace = None + if truesocket_recording_dir: + namespace = ".".join( + ( + instance.__class__.__module__, + instance.__class__.__name__, + test.__name__, + ) + ) + with cls( + instance, + namespace=namespace, + truesocket_recording_dir=truesocket_recording_dir, + ): + return test(*args, **kw) -mocketize = Mocketizer.wrap +if decorator.__version__ < '5': + mocketize = decorator.decorator(wrapper) +else: + mocketize = decorator.decorator(wrapper, kwsyntax=True)