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)
            else:
                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)
                    else:
                        return func(*args, **kwargs)
        raise NotConnectedError("not available connection")
    return wrapper
