YARN-5317. testAMRestartNotLostContainerCompleteMsg may fail. Contributed by sandflee

(cherry picked from commit 10b704c5946afe7bfd4a6be40192ce7ca745d817)

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/MockRM.java

@@ -741,7 +741,12 @@ public class MockRM extends ResourceManager {
     return client.failApplicationAttempt(req);
   }
 
-  // from AMLauncher
+  /**
+   * recommend to use launchAM, or use sendAMLaunched like:
+   * 1, wait RMAppAttempt scheduled
+   * 2, send node heartbeat
+   * 3, sendAMLaunched
+   */
   public MockAM sendAMLaunched(ApplicationAttemptId appAttemptId)
       throws Exception {
     MockAM am = new MockAM(getRMContext(), masterService, appAttemptId);

hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-resourcemanager/src/test/java/org/apache/hadoop/yarn/server/resourcemanager/applicationsmanager/TestAMRestart.java

@@ -149,12 +149,11 @@ public class TestAMRestart {
     Assert.assertFalse(newAttemptId.equals(am1.getApplicationAttemptId()));
 
     // launch the new AM
-    RMAppAttempt attempt2 = app1.getCurrentAppAttempt();
-    nm1.nodeHeartbeat(true);
-    MockAM am2 = rm1.sendAMLaunched(attempt2.getAppAttemptId());
+    MockAM am2 = rm1.launchAM(app1, rm1, nm1);
     RegisterApplicationMasterResponse registerResponse =
         am2.registerAppAttempt();
 
+
     // Assert two containers are running: container2 and container3;
     Assert.assertEquals(2, registerResponse.getContainersFromPreviousAttempts()
       .size());
@@ -837,10 +836,7 @@ public class TestAMRestart {
 
     // launch the new AM
     RMAppAttempt attempt2 = app1.getCurrentAppAttempt();
-    nm1.nodeHeartbeat(true);
-    MockAM am2 = rm1.sendAMLaunched(attempt2.getAppAttemptId());
-    am2.registerAppAttempt();
-    rm1.waitForState(app1.getApplicationId(), RMAppState.RUNNING);
+    MockAM am2 = rm1.launchAndRegisterAM(app1, rm1, nm1);
 
     // whether new AM could get container complete msg
     AllocateResponse allocateResponse = am2.allocate(