def METHOD_NAME(
    context_fn: _ContextFactory,
    environment_fn: Optional[_EnvironmentFactory] = None,
):
  """Returns a decorator for running a test in a context.
  Args:
    context_fn: A `Callable` that constructs a `tff.framework.AsyncContext` or
      `tff.framework.SyncContext` to install beore invoking the decorated
      function.
    environment_fn: A `Callable` that constructs a list of
      `contextlib.AbstractContextManager` to enter before invoking the decorated
      function.
  """
  def decorator(fn):
    @contextlib.contextmanager
    def install_context(
        context_fn: _ContextFactory,
        environment_fn: Optional[_EnvironmentFactory] = None,
    ):
      context = context_fn()
      with context_stack_impl.context_stack.install(context):
        if environment_fn is not None:
          with contextlib.ExitStack() as stack:
            context_managers = environment_fn()
            for context_manager in context_managers:
              stack.enter_context(context_manager)
            yield
        else:
          yield
    if asyncio.iscoroutinefunction(fn):
      @functools.wraps(fn)
      async def wrapper(*args, **kwargs):
        with install_context(context_fn, environment_fn):
          return await fn(*args, **kwargs)
    else:
      @functools.wraps(fn)
      def wrapper(*args, **kwargs):
        with install_context(context_fn, environment_fn):
          return fn(*args, **kwargs)
    return wrapper
  return decorator