from smb.base import NotConnectedError


def check_connection(func):
    def wrapper(*args, **kwargs):
        self = args[0]
        if self.cfg.master_node.current:
            if not self.__check_connection_node(self.cfg.master_node):
                if self.__reconnect_node(self.cfg.master_node):
                    return func(*args, **kwargs)
        if self.cfg.backup_node is not None:
            if self.cfg.master_node.current:
                if self.__reconnect_node(self.cfg.backup_node):
                    return func(*args, **kwargs)
            if self.cfg.backup_node.current:
                if self.__reconnect_node(self.cfg.master_node):
                    return func(*args, **kwargs)
                else:
                    if not self.__check_connection_node(self.cfg.backup_node):
                        if self.__reconnect_node(self.cfg.backup_node):
                            return func(*args, **kwargs)
        raise NotConnectedError("not available connection")
    return wrapper
