def with_wrapper()

in testslide/mock_callable.py [0:0]

• 20 lines of code
• 4 McCabe index (conditional complexity)

    def with_wrapper(self, func: Callable) -> "_MockAsyncCallableDSL":
        """
        Replace callable with given wrapper async function, that will be called as:

          await func(original_async_func, *args, **kwargs)

        receiving the original function as the first argument as well as any given
        arguments.
        """
        if not callable(func):
            raise ValueError("{} must be callable.".format(func))

        if not self._original_callable:
            raise ValueError("Can not wrap original callable that does not exist.")

        @functools.wraps(func)
        async def wrapper(*args: Any, **kwargs: Any) -> Any:
            coro = func(self._original_callable, *args, **kwargs)
            if not _is_coroutine(coro):
                raise NotACoroutine(
                    f"Function did not return a coroutine.\n"
                    f"{func} must return a coroutine."
                )
            return await coro

        self._add_runner(
            _AsyncImplementationRunner(
                self._original_target, self._method, self._original_callable, wrapper
            )
        )
        return self