YARN-10553. Refactor TestDistributedShell (#4159)

(cherry picked from commit 890f2da)

 Conflicts:
	hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/test/java/org/apache/hadoop/yarn/applications/distributedshell/TestDSWithMultipleNodeManager.java
	hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/test/java/org/apache/hadoop/yarn/applications/distributedshell/TestDistributedShell.java

Co-authored-by: Ahmed Hussein <[email protected]>

Author: Akira Ajisaka

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/main/java/org/apache/hadoop/yarn/applications/distributedshell/ApplicationMaster.java

@@ -781,7 +781,7 @@ private void printUsage(Options opts) {
     new HelpFormatter().printHelp("ApplicationMaster", opts);
   }
 
-  private void cleanup() {
+  protected void cleanup() {
     try {
       appSubmitterUgi.doAs(new PrivilegedExceptionAction<Void>() {
         @Override

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-applications-distributedshell/src/main/java/org/apache/hadoop/yarn/applications/distributedshell/Client.java

@@ -1414,21 +1414,19 @@ protected void sendStopSignal() {
     }
     int waitCount = 0;
     LOG.info("Waiting for Client to exit loop");
-    while (!isRunning.get()) {
+    while (isRunning.get()) {
       try {
         Thread.sleep(50);
       } catch (InterruptedException ie) {
         // do nothing
       } finally {
-        waitCount++;
-        if (isRunning.get() || waitCount > 2000) {
+        if (++waitCount > 2000) {
           break;
         }
       }
     }
-    LOG.info("Stopping yarnClient within the Client");
+    LOG.info("Stopping yarnClient within the DS Client");
     yarnClient.stop();
-    yarnClient.waitForServiceToStop(clientTimeout);
     LOG.info("done stopping Client");
   }
 }