YARN-6856. [YARN-3409] Support CLI for Node Attributes Mapping. Contributed by Naganarasimha G R.

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ha/HAAdmin.java

@@ -575,7 +575,7 @@ public abstract class HAAdmin extends Configured implements Tool {
     return 0;
   }
   
-  protected static class UsageInfo {
+  public static class UsageInfo {
     public final String args;
     public final String help;
     

hadoop-yarn-project/hadoop-yarn/bin/yarn

@@ -55,6 +55,7 @@ function hadoop_usage
   hadoop_add_subcommand "timelinereader" client "run the timeline reader server"
   hadoop_add_subcommand "timelineserver" daemon "run the timeline server"
   hadoop_add_subcommand "top" client "view cluster information"
+  hadoop_add_subcommand "node-attributes" "map node to attibutes"
   hadoop_add_subcommand "version" client "print the version"
   hadoop_generate_usage "${HADOOP_SHELL_EXECNAME}" true
 }
@@ -186,6 +187,10 @@ ${HADOOP_COMMON_HOME}/${HADOOP_COMMON_LIB_JARS_DIR}"
       hadoop_add_classpath "$HADOOP_YARN_HOME/$YARN_DIR/timelineservice/lib/*"
       HADOOP_CLASSNAME='org.apache.hadoop.yarn.server.timelineservice.reader.TimelineReaderServer'
     ;;
+	node-attributes)
+      HADOOP_SUBCMD_SUPPORTDAEMONIZATION="false"
+      HADOOP_CLASSNAME='org.apache.hadoop.yarn.client.cli.NodeAttributesCLI'
+	;;
     timelineserver)
       HADOOP_SUBCMD_SUPPORTDAEMONIZATION="true"
       HADOOP_CLASSNAME='org.apache.hadoop.yarn.server.applicationhistoryservice.ApplicationHistoryServer'

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java

@@ -0,0 +1,410 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.yarn.client.cli;
+
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import org.apache.commons.cli.CommandLine;
+import org.apache.commons.cli.GnuParser;
+import org.apache.commons.cli.MissingArgumentException;
+import org.apache.commons.cli.Options;
+import org.apache.commons.cli.ParseException;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.conf.Configured;
+import org.apache.hadoop.fs.CommonConfigurationKeys;
+import org.apache.hadoop.ha.HAAdmin.UsageInfo;
+import org.apache.hadoop.ipc.RemoteException;
+import org.apache.hadoop.util.Tool;
+import org.apache.hadoop.util.ToolRunner;
+import org.apache.hadoop.yarn.api.records.NodeAttribute;
+import org.apache.hadoop.yarn.api.records.NodeAttributeType;
+import org.apache.hadoop.yarn.client.ClientRMProxy;
+import org.apache.hadoop.yarn.conf.YarnConfiguration;
+import org.apache.hadoop.yarn.exceptions.YarnException;
+import org.apache.hadoop.yarn.server.api.ResourceManagerAdministrationProtocol;
+import org.apache.hadoop.yarn.server.api.protocolrecords.AttributeMappingOperationType;
+import org.apache.hadoop.yarn.server.api.protocolrecords.NodeToAttributes;
+import org.apache.hadoop.yarn.server.api.protocolrecords.NodesToAttributesMappingRequest;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
+
+/**
+ * CLI to map attributes to Nodes.
+ *
+ */
+public class NodeAttributesCLI extends Configured implements Tool {
+
+  protected static final String INVALID_MAPPING_ERR_MSG =
+      "Invalid Node to attribute mapping : ";
+
+  protected static final String USAGE_YARN_NODE_ATTRIBUTES =
+      "Usage: yarn node-attributes ";
+
+  protected static final String NO_MAPPING_ERR_MSG =
+      "No node-to-attributes mappings are specified";
+
+  protected final static Map<String, UsageInfo> NODE_ATTRIB_USAGE =
+      ImmutableMap.<String, UsageInfo>builder()
+          .put("-replace",
+              new UsageInfo(
+                  "<\"node1:attribute[(type)][=value],attribute1[=value],"
+                      + "attribute2  node2:attribute2[=value],attribute3\">",
+                  " Replace the node to attributes mapping information at the"
+                      + " ResourceManager with the new mapping. Currently"
+                      + " supported attribute type. And string is the default"
+                      + " type too. Attribute value if not specified for string"
+                      + " type value will be considered as empty string."
+                      + " Replaced node-attributes should not violate the"
+                      + " existing attribute to attribute type mapping."))
+          .put("-add",
+              new UsageInfo(
+                  "<\"node1:attribute[(type)][=value],attribute1[=value],"
+                      + "attribute2  node2:attribute2[=value],attribute3\">",
+                  " Adds or updates the node to attributes mapping information"
+                      + " at the ResourceManager. Currently supported attribute"
+                      + " type is string. And string is the default type too."
+                      + " Attribute value if not specified for string type"
+                      + " value will be considered as empty string. Added or"
+                      + " updated node-attributes should not violate the"
+                      + " existing attribute to attribute type mapping."))
+          .put("-remove",
+              new UsageInfo("<\"node1:attribute,attribute1 node2:attribute2\">",
+                  " Removes the specified node to attributes mapping"
+                      + " information at the ResourceManager"))
+          .put("-failOnUnknownNodes",
+              new UsageInfo("",
+                  "Can be used optionally along with other options. When its"
+                      + " set, it will fail if specified nodes are unknown."))
+          .build();
+
+  /** Output stream for errors, for use in tests. */
+  private PrintStream errOut = System.err;
+
+  public NodeAttributesCLI() {
+    super();
+  }
+
+  public NodeAttributesCLI(Configuration conf) {
+    super(conf);
+  }
+
+  protected void setErrOut(PrintStream errOut) {
+    this.errOut = errOut;
+  }
+
+  private void printHelpMsg(String cmd) {
+    StringBuilder builder = new StringBuilder();
+    UsageInfo usageInfo = null;
+    if (cmd != null && !(cmd.trim().isEmpty())) {
+      usageInfo = NODE_ATTRIB_USAGE.get(cmd);
+    }
+    if (usageInfo != null) {
+      if (usageInfo.args == null) {
+        builder.append("   " + cmd + ":\n" + usageInfo.help);
+      } else {
+        String space = (usageInfo.args == "") ? "" : " ";
+        builder.append(
+            "   " + cmd + space + usageInfo.args + " :\n" + usageInfo.help);
+      }
+    } else {
+      // help for all commands
+      builder.append("Usage: yarn node-attributes\n");
+      for (Map.Entry<String, UsageInfo> cmdEntry : NODE_ATTRIB_USAGE
+          .entrySet()) {
+        usageInfo = cmdEntry.getValue();
+        builder.append("   " + cmdEntry.getKey() + " " + usageInfo.args
+            + " :\n " + usageInfo.help + "\n");
+      }
+      builder.append("   -help" + " [cmd]\n");
+    }
+    errOut.println(builder);
+  }
+
+  private static void buildIndividualUsageMsg(String cmd,
+      StringBuilder builder) {
+    UsageInfo usageInfo = NODE_ATTRIB_USAGE.get(cmd);
+    if (usageInfo == null) {
+      return;
+    }
+    if (usageInfo.args == null) {
+      builder.append(USAGE_YARN_NODE_ATTRIBUTES + cmd + "\n");
+    } else {
+      String space = (usageInfo.args == "") ? "" : " ";
+      builder.append(
+          USAGE_YARN_NODE_ATTRIBUTES + cmd + space + usageInfo.args + "\n");
+    }
+  }
+
+  private static void buildUsageMsgForAllCmds(StringBuilder builder) {
+    builder.append("Usage: yarn node-attributes\n");
+    for (Map.Entry<String, UsageInfo> cmdEntry : NODE_ATTRIB_USAGE.entrySet()) {
+      UsageInfo usageInfo = cmdEntry.getValue();
+      builder.append("   " + cmdEntry.getKey() + " " + usageInfo.args + "\n");
+    }
+    builder.append("   -help" + " [cmd]\n");
+  }
+
+  /**
+   * Displays format of commands.
+   *
+   * @param cmd The command that is being executed.
+   */
+  private void printUsage(String cmd) {
+    StringBuilder usageBuilder = new StringBuilder();
+    if (NODE_ATTRIB_USAGE.containsKey(cmd)) {
+      buildIndividualUsageMsg(cmd, usageBuilder);
+    } else {
+      buildUsageMsgForAllCmds(usageBuilder);
+    }
+    errOut.println(usageBuilder);
+  }
+
+  private void printUsage() {
+    printUsage("");
+  }
+
+  protected ResourceManagerAdministrationProtocol createAdminProtocol()
+      throws IOException {
+    // Get the current configuration
+    final YarnConfiguration conf = new YarnConfiguration(getConf());
+    return ClientRMProxy.createRMProxy(conf,
+        ResourceManagerAdministrationProtocol.class);
+  }
+
+  @Override
+  public void setConf(Configuration conf) {
+    if (conf != null) {
+      conf = addSecurityConfiguration(conf);
+    }
+    super.setConf(conf);
+  }
+
+  /**
+   * Add the requisite security principal settings to the given Configuration,
+   * returning a copy.
+   *
+   * @param conf the original config
+   * @return a copy with the security settings added
+   */
+  private static Configuration addSecurityConfiguration(Configuration conf) {
+    // Make a copy so we don't mutate it. Also use an YarnConfiguration to
+    // force loading of yarn-site.xml.
+    conf = new YarnConfiguration(conf);
+    conf.set(CommonConfigurationKeys.HADOOP_SECURITY_SERVICE_USER_NAME_KEY,
+        conf.get(YarnConfiguration.RM_PRINCIPAL, ""));
+    return conf;
+  }
+
+  @Override
+  public int run(String[] args) throws Exception {
+    if (args.length < 1) {
+      printUsage();
+      return -1;
+    }
+
+    int exitCode = -1;
+    int i = 0;
+    String cmd = args[i++];
+
+    if ("-help".equals(cmd)) {
+      exitCode = 0;
+      if (args.length >= 2) {
+        printHelpMsg(args[i]);
+      } else {
+        printHelpMsg("");
+      }
+      return exitCode;
+    }
+
+    try {
+      if ("-replace".equals(cmd)) {
+        exitCode = handleNodeAttributeMapping(args,
+            AttributeMappingOperationType.REPLACE);
+      } else if ("-add".equals(cmd)) {
+        exitCode =
+            handleNodeAttributeMapping(args, AttributeMappingOperationType.ADD);
+      } else if ("-remove".equals(cmd)) {
+        exitCode = handleNodeAttributeMapping(args,
+            AttributeMappingOperationType.REMOVE);
+      } else {
+        exitCode = -1;
+        errOut.println(cmd.substring(1) + ": Unknown command");
+        printUsage();
+      }
+    } catch (IllegalArgumentException arge) {
+      exitCode = -1;
+      errOut.println(cmd.substring(1) + ": " + arge.getLocalizedMessage());
+      printUsage(cmd);
+    } catch (RemoteException e) {
+      //
+      // This is a error returned by hadoop server. Print
+      // out the first line of the error message, ignore the stack trace.
+      exitCode = -1;
+      try {
+        String[] content;
+        content = e.getLocalizedMessage().split("\n");
+        errOut.println(cmd.substring(1) + ": " + content[0]);
+      } catch (Exception ex) {
+        errOut.println(cmd.substring(1) + ": " + ex.getLocalizedMessage());
+      }
+    } catch (Exception e) {
+      exitCode = -1;
+      errOut.println(cmd.substring(1) + ": " + e.getLocalizedMessage());
+    }
+    return exitCode;
+  }
+
+  private int handleNodeAttributeMapping(String args[],
+      AttributeMappingOperationType operation)
+      throws IOException, YarnException, ParseException {
+    Options opts = new Options();
+    opts.addOption(operation.name().toLowerCase(), true,
+        operation.name().toLowerCase());
+    opts.addOption("failOnUnknownNodes", false, "Fail on unknown nodes.");
+    int exitCode = -1;
+    CommandLine cliParser = null;
+    try {
+      cliParser = new GnuParser().parse(opts, args);
+    } catch (MissingArgumentException ex) {
+      errOut.println(NO_MAPPING_ERR_MSG);
+      printUsage(args[0]);
+      return exitCode;
+    }
+    List<NodeToAttributes> buildNodeLabelsMapFromStr =
+        buildNodeLabelsMapFromStr(
+            cliParser.getOptionValue(operation.name().toLowerCase()),
+            operation != AttributeMappingOperationType.REPLACE, operation);
+    NodesToAttributesMappingRequest request = NodesToAttributesMappingRequest
+        .newInstance(operation, buildNodeLabelsMapFromStr,
+            cliParser.hasOption("failOnUnknownNodes"));
+    ResourceManagerAdministrationProtocol adminProtocol = createAdminProtocol();
+    adminProtocol.mapAttributesToNodes(request);
+    return 0;
+  }
+
+  /**
+   * args are expected to be of the format
+   * node1:java(string)=8,ssd(boolean)=false node2:ssd(boolean)=true
+   */
+  private List<NodeToAttributes> buildNodeLabelsMapFromStr(String args,
+      boolean validateForAttributes, AttributeMappingOperationType operation) {
+    List<NodeToAttributes> nodeToAttributesList = new ArrayList<>();
+    for (String nodeToAttributesStr : args.split("[ \n]")) {
+      // for each node to attribute mapping
+      nodeToAttributesStr = nodeToAttributesStr.trim();
+      if (nodeToAttributesStr.isEmpty()
+          || nodeToAttributesStr.startsWith("#")) {
+        continue;
+      }
+      if (nodeToAttributesStr.indexOf(":") == -1) {
+        throw new IllegalArgumentException(
+            INVALID_MAPPING_ERR_MSG + nodeToAttributesStr);
+      }
+      String[] nodeToAttributes = nodeToAttributesStr.split(":");
+      Preconditions.checkArgument(!nodeToAttributes[0].trim().isEmpty(),
+          "Node name cannot be empty");
+      String node = nodeToAttributes[0];
+      String[] attributeNameValueType = null;
+      List<NodeAttribute> attributesList = new ArrayList<>();
+      NodeAttributeType attributeType = NodeAttributeType.STRING;
+      String attributeValue;
+      String attributeName;
+      Set<String> attributeNamesMapped = new HashSet<>();
+
+      String attributesStr[];
+      if (nodeToAttributes.length == 2) {
+        // fetching multiple attributes for a node
+        attributesStr = nodeToAttributes[1].split(",");
+        for (String attributeStr : attributesStr) {
+          // get information about each attribute.
+          attributeNameValueType = attributeStr.split("="); // to find name
+                                                            // value
+          Preconditions.checkArgument(
+              !(attributeNameValueType[0] == null
+                  || attributeNameValueType[0].isEmpty()),
+              "Attribute name cannot be null or empty");
+          attributeValue = attributeNameValueType.length > 1
+              ? attributeNameValueType[1] : "";
+          int indexOfOpenBracket = attributeNameValueType[0].indexOf("(");
+          if (indexOfOpenBracket == -1) {
+            attributeName = attributeNameValueType[0];
+          } else if (indexOfOpenBracket == 0) {
+            throw new IllegalArgumentException("Attribute for node " + node
+                + " is not properly configured : " + attributeStr);
+          } else {
+            // attribute type has been explicitly configured
+            int indexOfCloseBracket = attributeNameValueType[0].indexOf(")");
+            if (indexOfCloseBracket == -1
+                || indexOfCloseBracket < indexOfOpenBracket) {
+              throw new IllegalArgumentException("Attribute for node " + node
+                  + " is not properly Configured : " + attributeStr);
+            }
+            String attributeTypeStr;
+            attributeName =
+                attributeNameValueType[0].substring(0, indexOfOpenBracket);
+            attributeTypeStr = attributeNameValueType[0]
+                .substring(indexOfOpenBracket + 1, indexOfCloseBracket);
+            try {
+              attributeType = NodeAttributeType
+                  .valueOf(attributeTypeStr.trim().toUpperCase());
+            } catch (IllegalArgumentException e) {
+              throw new IllegalArgumentException(
+                  "Invalid Attribute type configuration : " + attributeTypeStr
+                      + " in " + attributeStr);
+            }
+          }
+          if (attributeNamesMapped.contains(attributeName)) {
+            throw new IllegalArgumentException("Attribute " + attributeName
+                + " has been mapped more than once in  : "
+                + nodeToAttributesStr);
+          }
+          // TODO when we support different type of attribute type we need to
+          // cross verify whether input attributes itself is not violating
+          // attribute Name to Type mapping.
+          attributesList.add(NodeAttribute.newInstance(attributeName.trim(),
+              attributeType, attributeValue.trim()));
+        }
+      }
+      if (validateForAttributes) {
+        Preconditions.checkArgument((attributesList.size() > 0),
+            "Attributes cannot be null or empty for Operation "
+                + operation.name() + " on the node " + node);
+      }
+      nodeToAttributesList
+          .add(NodeToAttributes.newInstance(node, attributesList));
+    }
+
+    if (nodeToAttributesList.isEmpty()) {
+      throw new IllegalArgumentException(NO_MAPPING_ERR_MSG);
+    }
+    return nodeToAttributesList;
+  }
+
+  public static void main(String[] args) throws Exception {
+    int result = ToolRunner.run(new NodeAttributesCLI(), args);
+    System.exit(result);
+  }
+}

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java

@@ -0,0 +1,328 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hadoop.yarn.client.cli;
+
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.PrintStream;
+import java.util.ArrayList;
+import java.util.List;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.yarn.api.records.NodeAttribute;
+import org.apache.hadoop.yarn.api.records.NodeAttributeType;
+import org.apache.hadoop.yarn.exceptions.YarnException;
+import org.apache.hadoop.yarn.server.api.ResourceManagerAdministrationProtocol;
+import org.apache.hadoop.yarn.server.api.protocolrecords.AttributeMappingOperationType;
+import org.apache.hadoop.yarn.server.api.protocolrecords.NodeToAttributes;
+import org.apache.hadoop.yarn.server.api.protocolrecords.NodesToAttributesMappingRequest;
+import org.apache.hadoop.yarn.server.api.protocolrecords.NodesToAttributesMappingResponse;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.google.common.base.Charsets;
+import com.google.common.base.Joiner;
+
+/**
+ * Test class for TestNodeAttributesCLI.
+ */
+public class TestNodeAttributesCLI {
+  private static final Logger LOG =
+      LoggerFactory.getLogger(TestNodeAttributesCLI.class);
+  private ResourceManagerAdministrationProtocol admin;
+  private NodesToAttributesMappingRequest request;
+  private NodeAttributesCLI nodeAttributesCLI;
+  private ByteArrayOutputStream errOutBytes = new ByteArrayOutputStream();
+  private String errOutput;
+
+  @Before
+  public void configure() throws IOException, YarnException {
+    admin = mock(ResourceManagerAdministrationProtocol.class);
+
+    when(admin.mapAttributesToNodes(any(NodesToAttributesMappingRequest.class)))
+        .thenAnswer(new Answer<NodesToAttributesMappingResponse>() {
+          @Override
+          public NodesToAttributesMappingResponse answer(
+              InvocationOnMock invocation) throws Throwable {
+            request =
+                (NodesToAttributesMappingRequest) invocation.getArguments()[0];
+            return NodesToAttributesMappingResponse.newInstance();
+          }
+        });
+
+    nodeAttributesCLI = new NodeAttributesCLI(new Configuration()) {
+      @Override
+      protected ResourceManagerAdministrationProtocol createAdminProtocol()
+          throws IOException {
+        return admin;
+      }
+    };
+
+    nodeAttributesCLI.setErrOut(new PrintStream(errOutBytes));
+  }
+
+  @Test
+  public void testHelp() throws Exception {
+    String[] args = new String[] { "-help", "-replace" };
+    assertTrue("It should have succeeded help for replace", 0 == runTool(args));
+    assertOutputContains(
+        "-replace <\"node1:attribute[(type)][=value],attribute1"
+            + "[=value],attribute2  node2:attribute2[=value],attribute3\"> :");
+    assertOutputContains("Replace the node to attributes mapping information at"
+        + " the ResourceManager with the new mapping. Currently supported"
+        + " attribute type. And string is the default type too. Attribute value"
+        + " if not specified for string type value will be considered as empty"
+        + " string. Replaced node-attributes should not violate the existing"
+        + " attribute to attribute type mapping.");
+
+    args = new String[] { "-help", "-remove" };
+    assertTrue("It should have succeeded help for replace", 0 == runTool(args));
+    assertOutputContains(
+        "-remove <\"node1:attribute,attribute1" + " node2:attribute2\"> :");
+    assertOutputContains("Removes the specified node to attributes mapping"
+        + " information at the ResourceManager");
+
+    args = new String[] { "-help", "-add" };
+    assertTrue("It should have succeeded help for replace", 0 == runTool(args));
+    assertOutputContains("-add <\"node1:attribute[(type)][=value],"
+        + "attribute1[=value],attribute2  node2:attribute2[=value],attribute3\">"
+        + " :");
+    assertOutputContains("Adds or updates the node to attributes mapping"
+        + " information at the ResourceManager. Currently supported attribute"
+        + " type is string. And string is the default type too. Attribute value"
+        + " if not specified for string type value will be considered as empty"
+        + " string. Added or updated node-attributes should not violate the"
+        + " existing attribute to attribute type mapping.");
+
+    args = new String[] { "-help", "-failOnUnknownNodes" };
+    assertTrue("It should have succeeded help for replace", 0 == runTool(args));
+    assertOutputContains("-failOnUnknownNodes :");
+    assertOutputContains("Can be used optionally along with other options. When"
+        + " its set, it will fail if specified nodes are unknown.");
+  }
+
+  @Test
+  public void testReplace() throws Exception {
+    // --------------------------------
+    // failure scenarios
+    // --------------------------------
+    // parenthesis not match
+    String[] args = new String[] { "-replace", "x(" };
+    assertTrue("It should have failed as no node is specified",
+        0 != runTool(args));
+    assertFailureMessageContains(NodeAttributesCLI.INVALID_MAPPING_ERR_MSG);
+
+    // parenthesis not match
+    args = new String[] { "-replace", "x:(=abc" };
+    assertTrue(
+        "It should have failed as no closing parenthesis is not specified",
+        0 != runTool(args));
+    assertFailureMessageContains(
+        "Attribute for node x is not properly configured : (=abc");
+
+    args = new String[] { "-replace", "x:()=abc" };
+    assertTrue("It should have failed as no type specified inside parenthesis",
+        0 != runTool(args));
+    assertFailureMessageContains(
+        "Attribute for node x is not properly configured : ()=abc");
+
+    args = new String[] { "-replace", ":x(string)" };
+    assertTrue("It should have failed as no node is specified",
+        0 != runTool(args));
+    assertFailureMessageContains("Node name cannot be empty");
+
+    // Not expected key=value specifying inner parenthesis
+    args = new String[] { "-replace", "x:(key=value)" };
+    assertTrue(0 != runTool(args));
+    assertFailureMessageContains(
+        "Attribute for node x is not properly configured : (key=value)");
+
+    // Should fail as no attributes specified
+    args = new String[] { "-replace" };
+    assertTrue("Should fail as no attribute mappings specified",
+        0 != runTool(args));
+    assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG);
+
+    // no labels, should fail
+    args = new String[] { "-replace", "-failOnUnknownNodes",
+        "x:key(string)=value,key2=val2" };
+    assertTrue("Should fail as no attribute mappings specified for replace",
+        0 != runTool(args));
+    assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG);
+
+    // no labels, should fail
+    args = new String[] { "-replace", " " };
+    assertTrue(0 != runTool(args));
+    assertFailureMessageContains(NodeAttributesCLI.NO_MAPPING_ERR_MSG);
+
+    args = new String[] { "-replace", ", " };
+    assertTrue(0 != runTool(args));
+    assertFailureMessageContains(NodeAttributesCLI.INVALID_MAPPING_ERR_MSG);
+    // --------------------------------
+    // success scenarios
+    // --------------------------------
+    args = new String[] { "-replace",
+        "x:key(string)=value,key2=val2 y:key2=val23,key3 z:key4" };
+    assertTrue("Should not fail as attribute has been properly mapped",
+        0 == runTool(args));
+    List<NodeToAttributes> nodeAttributesList = new ArrayList<>();
+    List<NodeAttribute> attributes = new ArrayList<>();
+    attributes.add(
+        NodeAttribute.newInstance("key", NodeAttributeType.STRING, "value"));
+    attributes.add(
+        NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "val2"));
+    nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes));
+
+    // for node y
+    attributes = new ArrayList<>();
+    attributes.add(
+        NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "val23"));
+    attributes
+        .add(NodeAttribute.newInstance("key3", NodeAttributeType.STRING, ""));
+    nodeAttributesList.add(NodeToAttributes.newInstance("y", attributes));
+
+    // for node y
+    attributes = new ArrayList<>();
+    attributes.add(
+        NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "val23"));
+    attributes
+        .add(NodeAttribute.newInstance("key3", NodeAttributeType.STRING, ""));
+    nodeAttributesList.add(NodeToAttributes.newInstance("y", attributes));
+
+    // for node z
+    attributes = new ArrayList<>();
+    attributes
+        .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, ""));
+    nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes));
+
+    NodesToAttributesMappingRequest expected =
+        NodesToAttributesMappingRequest.newInstance(
+            AttributeMappingOperationType.REPLACE, nodeAttributesList, false);
+    assertTrue(request.equals(expected));
+  }
+
+  @Test
+  public void testRemove() throws Exception {
+    // --------------------------------
+    // failure scenarios
+    // --------------------------------
+    // parenthesis not match
+    String[] args = new String[] { "-remove", "x:" };
+    assertTrue("It should have failed as no node is specified",
+        0 != runTool(args));
+    assertFailureMessageContains(
+        "Attributes cannot be null or empty for Operation REMOVE on the node x");
+    // --------------------------------
+    // success scenarios
+    // --------------------------------
+    args =
+        new String[] { "-remove", "x:key2,key3 z:key4", "-failOnUnknownNodes" };
+    assertTrue("Should not fail as attribute has been properly mapped",
+        0 == runTool(args));
+    List<NodeToAttributes> nodeAttributesList = new ArrayList<>();
+    List<NodeAttribute> attributes = new ArrayList<>();
+    attributes
+        .add(NodeAttribute.newInstance("key2", NodeAttributeType.STRING, ""));
+    attributes
+        .add(NodeAttribute.newInstance("key3", NodeAttributeType.STRING, ""));
+    nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes));
+
+    // for node z
+    attributes = new ArrayList<>();
+    attributes
+        .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, ""));
+    nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes));
+
+    NodesToAttributesMappingRequest expected =
+        NodesToAttributesMappingRequest.newInstance(
+            AttributeMappingOperationType.REMOVE, nodeAttributesList, true);
+    assertTrue(request.equals(expected));
+  }
+
+  @Test
+  public void testAdd() throws Exception {
+    // --------------------------------
+    // failure scenarios
+    // --------------------------------
+    // parenthesis not match
+    String[] args = new String[] { "-add", "x:" };
+    assertTrue("It should have failed as no node is specified",
+        0 != runTool(args));
+    assertFailureMessageContains(
+        "Attributes cannot be null or empty for Operation ADD on the node x");
+    // --------------------------------
+    // success scenarios
+    // --------------------------------
+    args = new String[] { "-add", "x:key2=123,key3=abc z:key4(string)",
+        "-failOnUnknownNodes" };
+    assertTrue("Should not fail as attribute has been properly mapped",
+        0 == runTool(args));
+    List<NodeToAttributes> nodeAttributesList = new ArrayList<>();
+    List<NodeAttribute> attributes = new ArrayList<>();
+    attributes.add(
+        NodeAttribute.newInstance("key2", NodeAttributeType.STRING, "123"));
+    attributes.add(
+        NodeAttribute.newInstance("key3", NodeAttributeType.STRING, "abc"));
+    nodeAttributesList.add(NodeToAttributes.newInstance("x", attributes));
+
+    // for node z
+    attributes = new ArrayList<>();
+    attributes
+        .add(NodeAttribute.newInstance("key4", NodeAttributeType.STRING, ""));
+    nodeAttributesList.add(NodeToAttributes.newInstance("z", attributes));
+
+    NodesToAttributesMappingRequest expected =
+        NodesToAttributesMappingRequest.newInstance(
+            AttributeMappingOperationType.ADD, nodeAttributesList, true);
+    assertTrue(request.equals(expected));
+  }
+
+  private void assertFailureMessageContains(String... messages) {
+    assertOutputContains(messages);
+    assertOutputContains(NodeAttributesCLI.USAGE_YARN_NODE_ATTRIBUTES);
+  }
+
+  private void assertOutputContains(String... messages) {
+    for (String message : messages) {
+      if (!errOutput.contains(message)) {
+        fail("Expected output to contain '" + message
+            + "' but err_output was:\n" + errOutput);
+      }
+    }
+  }
+
+  private int runTool(String... args) throws Exception {
+    errOutBytes.reset();
+    LOG.info("Running: NodeAttributesCLI " + Joiner.on(" ").join(args));
+    int ret = nodeAttributesCLI.run(args);
+    errOutput = new String(errOutBytes.toByteArray(), Charsets.UTF_8);
+    LOG.info("Err_output:\n" + errOutput);
+    return ret;
+  }
+}