ClusterCLI extends YarnCLI { @Override public int run(String[] args) throws Exception { Options opts = new Options(); opts.addOption("lnl", LIST_LABELS_CMD, false, "List cluster node-label collection"); opts.addOption("h", HELP_CMD, false, "Displays help for all commands."); opts.addOption("dnl", DIRECTLY_ACCESS_NODE_LABEL_STORE, false, "Directly access node label store, " + "with this option, all node label related operations" + " will NOT connect RM. Instead, they will" + " access/modify stored node labels directly." + " By default, it is false (access via RM)." + " AND PLEASE NOTE: if you configured " + YarnConfiguration.FS_NODE_LABELS_STORE_ROOT_DIR + " to a local directory" + " (instead of NFS or HDFS), this option will only work" + " when the command run on the machine where RM is running." + " Also, this option is UNSTABLE, could be removed in future" + " releases."); int exitCode = -1; CommandLine parsedCli = null; try { parsedCli = new GnuParser().parse(opts, args); } catch (MissingArgumentException ex) { sysout.println("Missing argument for options"); printUsage(opts); return exitCode; } if (parsedCli.hasOption(DIRECTLY_ACCESS_NODE_LABEL_STORE)) { accessLocal = true; } if (parsedCli.hasOption(LIST_LABELS_CMD)) { printClusterNodeLabels(); } else if (parsedCli.hasOption(HELP_CMD)) { printUsage(opts); return 0; } else { syserr.println("Invalid Command Usage : "); printUsage(opts); } return 0; }  static void main(String[] args); @Override int run(String[] args); static final String LIST_LABELS_CMD; static final String DIRECTLY_ACCESS_NODE_LABEL_STORE; static final String CMD; }
@Test public void testGetClusterNodeLabels() throws Exception { YarnClient client = mock(YarnClient.class); when(client.getClusterNodeLabels()).thenReturn( ImmutableSet.of("label1", "label2")); ClusterCLI cli = new ClusterCLI(); cli.setClient(client); cli.setSysOutPrintStream(sysOut); cli.setSysErrPrintStream(sysErr); int rc = cli.run(new String[] { ClusterCLI.CMD, "-" + ClusterCLI.LIST_LABELS_CMD }); assertEquals(0, rc); ByteArrayOutputStream baos = new ByteArrayOutputStream(); PrintWriter pw = new PrintWriter(baos); pw.print("Node Labels: label1,label2"); pw.close(); verify(sysOut).println(baos.toString("UTF-8")); }
@Test public void testGetClusterNodeLabelsWithLocalAccess() throws Exception { YarnClient client = mock(YarnClient.class); when(client.getClusterNodeLabels()).thenReturn( ImmutableSet.of("remote1", "remote2")); ClusterCLI cli = new ClusterCLI(); cli.setClient(client); cli.setSysOutPrintStream(sysOut); cli.setSysErrPrintStream(sysErr); ClusterCLI.localNodeLabelsManager = mock(CommonNodeLabelsManager.class); when(ClusterCLI.localNodeLabelsManager.getClusterNodeLabels()) .thenReturn(ImmutableSet.of("local1", "local2")); int rc = cli.run(new String[] { ClusterCLI.CMD, "-" + ClusterCLI.LIST_LABELS_CMD, "-" + ClusterCLI.DIRECTLY_ACCESS_NODE_LABEL_STORE }); assertEquals(0, rc); ByteArrayOutputStream baos = new ByteArrayOutputStream(); PrintWriter pw = new PrintWriter(baos); pw.print("Node Labels: local1,local2"); pw.close(); verify(sysOut).println(baos.toString("UTF-8")); }
@Test public void testGetEmptyClusterNodeLabels() throws Exception { YarnClient client = mock(YarnClient.class); when(client.getClusterNodeLabels()).thenReturn(new HashSet<String>()); ClusterCLI cli = new ClusterCLI(); cli.setClient(client); cli.setSysOutPrintStream(sysOut); cli.setSysErrPrintStream(sysErr); int rc = cli.run(new String[] { ClusterCLI.CMD, "-" + ClusterCLI.LIST_LABELS_CMD }); assertEquals(0, rc); ByteArrayOutputStream baos = new ByteArrayOutputStream(); PrintWriter pw = new PrintWriter(baos); pw.print("Node Labels: "); pw.close(); verify(sysOut).println(baos.toString("UTF-8")); }
@Test public void testHelp() throws Exception { ClusterCLI cli = new ClusterCLI(); cli.setSysOutPrintStream(sysOut); cli.setSysErrPrintStream(sysErr); int rc = cli.run(new String[] { "cluster", "--help" }); assertEquals(0, rc); ByteArrayOutputStream baos = new ByteArrayOutputStream(); PrintWriter pw = new PrintWriter(baos); pw.println("usage: yarn cluster"); pw.println(" -dnl,--directly-access-node-label-store Directly access node label"); pw.println(" store, with this option, all"); pw.println(" node label related operations"); pw.println(" will NOT connect RM. Instead,"); pw.println(" they will access/modify stored"); pw.println(" node labels directly. By"); pw.println(" default, it is false (access"); pw.println(" via RM). AND PLEASE NOTE: if"); pw.println(" you configured"); pw.println(" yarn.node-labels.fs-store.root-"); pw.println(" dir to a local directory"); pw.println(" (instead of NFS or HDFS), this"); pw.println(" option will only work when the"); pw.println(" command run on the machine"); pw.println(" where RM is running. Also, this"); pw.println(" option is UNSTABLE, could be"); pw.println(" removed in future releases."); pw.println(" -h,--help Displays help for all commands."); pw.println(" -lnl,--list-node-labels List cluster node-label"); pw.println(" collection"); pw.close(); verify(sysOut).println(baos.toString("UTF-8")); }