import functools
from invoke.exceptions import UnexpectedExit


class RemoteError(Exception):
    pass


# wrap functions with remote connection calls in this to avoid cryptic errors
def wrap_canesm_remotefail(func, message=None):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        canesm_exception = None
        try:
            return func(*args, **kwargs)
        except UnexpectedExit as e:
            # Store the exception, but don't raise here or else we will get funny chained exceptions
            canesm_exception = e

        raise RemoteError(canesm_exception.result)
    return wrapper