YARN-2649. Fixed TestAMRMRPCNodeUpdates test failure. Contributed by Ming Ma

hadoop-yarn-project/CHANGES.txt

@@ -588,6 +588,8 @@ Release 2.6.0 - UNRELEASED
     YARN-1857. CapacityScheduler headroom doesn't account for other AM's running.
     (Chen He and Craig Welch via jianhe)
 
+    YARN-2649. Fixed TestAMRMRPCNodeUpdates test failure. (Ming Ma via jianhe)
+
 Release 2.5.1 - 2014-09-05
 
   INCOMPATIBLE CHANGES

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

@@ -386,7 +386,15 @@ public class MockRM extends ResourceManager {
     if (waitForAccepted) {
       waitForState(appId, RMAppState.ACCEPTED);
     }
-    return getRMContext().getRMApps().get(appId);
+    RMApp rmApp = getRMContext().getRMApps().get(appId);
+
+    // unmanaged AM won't go to RMAppAttemptState.SCHEDULED.
+    if (waitForAccepted && !unmanaged) {
+      waitForState(rmApp.getCurrentAppAttempt().getAppAttemptId(),
+          RMAppAttemptState.SCHEDULED);
+    }
+
+    return rmApp;
   }
 
   public MockNM registerNode(String nodeIdStr, int memory) throws Exception {