YARN-9995. Code cleanup in TestSchedConfCLI. Contributed by Bilwa S T.

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

@@ -65,6 +65,7 @@ import javax.servlet.http.HttpServletRequestWrapper;
 import javax.servlet.http.HttpServletResponse;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
@@ -73,17 +74,10 @@ import static org.junit.Assert.assertTrue;
  */
 public class TestSchedConfCLI extends JerseyTestBase {
 
-  private ByteArrayOutputStream sysOutStream;
-  private PrintStream sysOut;
-
-  private ByteArrayOutputStream sysErrStream;
-  private PrintStream sysErr;
-
   private SchedConfCLI cli;
 
   private static MockRM rm;
   private static String userName;
-  private static CapacitySchedulerConfiguration csConf;
 
   private static final File CONF_FILE = new File(new File("target",
       "test-classes"), YarnConfiguration.CS_CONFIGURATION_FILE);
@@ -100,14 +94,6 @@ public class TestSchedConfCLI extends JerseyTestBase {
 
   @Before
   public void setUp() {
-    sysOutStream = new ByteArrayOutputStream();
-    sysOut =  new PrintStream(sysOutStream);
-    System.setOut(sysOut);
-
-    sysErrStream = new ByteArrayOutputStream();
-    sysErr = new PrintStream(sysErrStream);
-    System.setErr(sysErr);
-
     cli = new SchedConfCLI();
   }
 
@@ -130,8 +116,8 @@ public class TestSchedConfCLI extends JerseyTestBase {
             + ioe.getMessage(), ioe);
       }
 
-      csConf = new CapacitySchedulerConfiguration(new Configuration(false),
+      CapacitySchedulerConfiguration csConf = new
-          false);
+          CapacitySchedulerConfiguration(new Configuration(false), false);
       setupQueueConfiguration(csConf);
 
       try {
@@ -245,51 +231,46 @@ public class TestSchedConfCLI extends JerseyTestBase {
 
   @Test(timeout = 10000)
   public void testInvalidConf() throws Exception {
-    // conf pair with no key should be invalid
+    ByteArrayOutputStream sysErrStream = new ByteArrayOutputStream();
-    int exitCode = cli.run(new String[] {"-add", "root.a:=confVal"});
+    PrintStream sysErr = new PrintStream(sysErrStream);
-    assertTrue("Should return an error code", exitCode != 0);
+    System.setErr(sysErr);
-    assertTrue(sysErrStream.toString().contains("Specify configuration key " +
-        "value as confKey=confVal."));
-    exitCode = cli.run(new String[] {"-update", "root.a:=confVal"});
-    assertTrue("Should return an error code", exitCode != 0);
-    assertTrue(sysErrStream.toString().contains("Specify configuration key " +
-        "value as confKey=confVal."));
 
-    exitCode = cli.run(new String[] {"-add", "root.a:confKey=confVal=conf"});
+    // conf pair with no key should be invalid
-    assertTrue("Should return an error code", exitCode != 0);
+    executeCommand(sysErrStream, "-add", "root.a:=confVal");
-    assertTrue(sysErrStream.toString().contains("Specify configuration key " +
+    executeCommand(sysErrStream, "-update", "root.a:=confVal");
-        "value as confKey=confVal."));
+    executeCommand(sysErrStream, "-add", "root.a:confKey=confVal=conf");
-    exitCode = cli.run(new String[] {"-update", "root.a:confKey=confVal=c"});
+    executeCommand(sysErrStream, "-update", "root.a:confKey=confVal=c");
-    assertTrue("Should return an error code", exitCode != 0);
+  }
-    assertTrue(sysErrStream.toString().contains("Specify configuration key " +
+
-        "value as confKey=confVal."));
+  private void executeCommand(ByteArrayOutputStream sysErrStream, String op,
+      String queueConf) throws Exception {
+    int exitCode = cli.run(new String[] {op, queueConf});
+    assertNotEquals("Should return an error code", 0, exitCode);
+    assertTrue(sysErrStream.toString()
+        .contains("Specify configuration key " + "value as confKey=confVal."));
   }
 
   @Test(timeout = 10000)
   public void testAddQueues() {
     SchedConfUpdateInfo schedUpdateInfo = new SchedConfUpdateInfo();
     cli.addQueues("root.a:a1=aVal1,a2=aVal2,a3=", schedUpdateInfo);
-    QueueConfigInfo addInfo = schedUpdateInfo.getAddQueueInfo().get(0);
+    Map<String, String> paramValues = new HashMap<>();
-    assertEquals("root.a", addInfo.getQueue());
+    List<QueueConfigInfo> addQueueInfo = schedUpdateInfo.getAddQueueInfo();
-    Map<String, String> params = addInfo.getParams();
+    paramValues.put("a1", "aVal1");
-    assertEquals(3, params.size());
+    paramValues.put("a2", "aVal2");
-    assertEquals("aVal1", params.get("a1"));
+    paramValues.put("a3", null);
-    assertEquals("aVal2", params.get("a2"));
+    validateQueueConfigInfo(addQueueInfo, 0, "root.a", paramValues);
-    assertNull(params.get("a3"));
 
     schedUpdateInfo = new SchedConfUpdateInfo();
     cli.addQueues("root.b:b1=bVal1;root.c:c1=cVal1", schedUpdateInfo);
-    assertEquals(2, schedUpdateInfo.getAddQueueInfo().size());
+    addQueueInfo = schedUpdateInfo.getAddQueueInfo();
-    QueueConfigInfo bAddInfo = schedUpdateInfo.getAddQueueInfo().get(0);
+    assertEquals(2, addQueueInfo.size());
-    assertEquals("root.b", bAddInfo.getQueue());
+    paramValues.clear();
-    Map<String, String> bParams = bAddInfo.getParams();
+    paramValues.put("b1", "bVal1");
-    assertEquals(1, bParams.size());
+    validateQueueConfigInfo(addQueueInfo, 0, "root.b", paramValues);
-    assertEquals("bVal1", bParams.get("b1"));
+    paramValues.clear();
-    QueueConfigInfo cAddInfo = schedUpdateInfo.getAddQueueInfo().get(1);
+    paramValues.put("c1", "cVal1");
-    assertEquals("root.c", cAddInfo.getQueue());
+    validateQueueConfigInfo(addQueueInfo, 1, "root.c", paramValues);
-    Map<String, String> cParams = cAddInfo.getParams();
-    assertEquals(1, cParams.size());
-    assertEquals("cVal1", cParams.get("c1"));
   }
 
   @Test(timeout = 10000)
@@ -306,28 +287,35 @@ public class TestSchedConfCLI extends JerseyTestBase {
   @Test(timeout = 10000)
   public void testUpdateQueues() {
     SchedConfUpdateInfo schedUpdateInfo = new SchedConfUpdateInfo();
+    Map<String, String> paramValues = new HashMap<>();
     cli.updateQueues("root.a:a1=aVal1,a2=aVal2,a3=", schedUpdateInfo);
-    QueueConfigInfo updateInfo = schedUpdateInfo.getUpdateQueueInfo().get(0);
+    List<QueueConfigInfo> updateQueueInfo = schedUpdateInfo
-    assertEquals("root.a", updateInfo.getQueue());
+        .getUpdateQueueInfo();
-    Map<String, String> params = updateInfo.getParams();
+    paramValues.put("a1", "aVal1");
-    assertEquals(3, params.size());
+    paramValues.put("a2", "aVal2");
-    assertEquals("aVal1", params.get("a1"));
+    paramValues.put("a3", null);
-    assertEquals("aVal2", params.get("a2"));
+    validateQueueConfigInfo(updateQueueInfo, 0, "root.a", paramValues);
-    assertNull(params.get("a3"));
 
     schedUpdateInfo = new SchedConfUpdateInfo();
     cli.updateQueues("root.b:b1=bVal1;root.c:c1=cVal1", schedUpdateInfo);
-    assertEquals(2, schedUpdateInfo.getUpdateQueueInfo().size());
+    updateQueueInfo = schedUpdateInfo.getUpdateQueueInfo();
-    QueueConfigInfo bUpdateInfo = schedUpdateInfo.getUpdateQueueInfo().get(0);
+    assertEquals(2, updateQueueInfo.size());
-    assertEquals("root.b", bUpdateInfo.getQueue());
+    paramValues.clear();
-    Map<String, String> bParams = bUpdateInfo.getParams();
+    paramValues.put("b1", "bVal1");
-    assertEquals(1, bParams.size());
+    validateQueueConfigInfo(updateQueueInfo, 0, "root.b", paramValues);
-    assertEquals("bVal1", bParams.get("b1"));
+    paramValues.clear();
-    QueueConfigInfo cUpdateInfo = schedUpdateInfo.getUpdateQueueInfo().get(1);
+    paramValues.put("c1", "cVal1");
-    assertEquals("root.c", cUpdateInfo.getQueue());
+    validateQueueConfigInfo(updateQueueInfo, 1, "root.c", paramValues);
-    Map<String, String> cParams = cUpdateInfo.getParams();
+  }
-    assertEquals(1, cParams.size());
+
-    assertEquals("cVal1", cParams.get("c1"));
+  private void validateQueueConfigInfo(
+      List<QueueConfigInfo> updateQueueInfo, int index, String queuename,
+      Map<String, String> paramValues) {
+    QueueConfigInfo updateInfo = updateQueueInfo.get(index);
+    assertEquals(queuename, updateInfo.getQueue());
+    Map<String, String> params = updateInfo.getParams();
+    assertEquals(paramValues.size(), params.size());
+    paramValues.forEach((k, v) -> assertEquals(v, params.get(k)));
   }
 
   @Test(timeout = 10000)
@@ -335,9 +323,17 @@ public class TestSchedConfCLI extends JerseyTestBase {
     SchedConfUpdateInfo schedUpdateInfo = new SchedConfUpdateInfo();
     cli.globalUpdates("schedKey1=schedVal1,schedKey2=schedVal2",
         schedUpdateInfo);
-    Map<String, String> globalInfo = schedUpdateInfo.getGlobalParams();
+    Map<String, String> paramValues = new HashMap<>();
-    assertEquals(2, globalInfo.size());
+    paramValues.put("schedKey1", "schedVal1");
-    assertEquals("schedVal1", globalInfo.get("schedKey1"));
+    paramValues.put("schedKey2", "schedVal2");
-    assertEquals("schedVal2", globalInfo.get("schedKey2"));
+    validateGlobalParams(schedUpdateInfo, paramValues);
   }
+
+  private void validateGlobalParams(SchedConfUpdateInfo schedUpdateInfo,
+      Map<String, String> paramValues) {
+    Map<String, String> globalInfo = schedUpdateInfo.getGlobalParams();
+    assertEquals(paramValues.size(), globalInfo.size());
+    paramValues.forEach((k, v) -> assertEquals(v, globalInfo.get(k)));
+  }
+
 }