diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ha/HAAdmin.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ha/HAAdmin.java index 9b7d7ba5d1a..8c92bd04471 100644 --- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/ha/HAAdmin.java +++ b/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; diff --git a/hadoop-yarn-project/hadoop-yarn/bin/yarn b/hadoop-yarn-project/hadoop-yarn/bin/yarn index 69afe6f88a7..7cd838fb813 100755 --- a/hadoop-yarn-project/hadoop-yarn/bin/yarn +++ b/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' diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/main/java/org/apache/hadoop/yarn/client/cli/NodeAttributesCLI.java new file mode 100644 index 00000000000..2eff155ee96 --- /dev/null +++ b/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 NODE_ATTRIB_USAGE = + ImmutableMap.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 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 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 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 buildNodeLabelsMapFromStr(String args, + boolean validateForAttributes, AttributeMappingOperationType operation) { + List 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 attributesList = new ArrayList<>(); + NodeAttributeType attributeType = NodeAttributeType.STRING; + String attributeValue; + String attributeName; + Set 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); + } +} diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-client/src/test/java/org/apache/hadoop/yarn/client/cli/TestNodeAttributesCLI.java new file mode 100644 index 00000000000..cc92a93eaf8 --- /dev/null +++ b/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() { + @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 nodeAttributesList = new ArrayList<>(); + List 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 nodeAttributesList = new ArrayList<>(); + List 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 nodeAttributesList = new ArrayList<>(); + List 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; + } +}