diff --git hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-tests/src/test/java/org/apache/hadoop/yarn/server/TestContainerManagerSecurity.java hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-tests/src/test/java/org/apache/hadoop/yarn/server/TestContainerManagerSecurity.java index 1cbad707808..bfd04b73e16 100644 --- hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-tests/src/test/java/org/apache/hadoop/yarn/server/TestContainerManagerSecurity.java +++ hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-tests/src/test/java/org/apache/hadoop/yarn/server/TestContainerManagerSecurity.java @@ -408,21 +408,18 @@ private void waitForContainerToFinishOnNM(ContainerId containerId) { Context nmContext = yarnCluster.getNodeManager(0).getNMContext(); int interval = 4 * 60; // Max time for container token to expire. - Assert.assertNotNull(nmContext.getContainers().containsKey(containerId)); - - // Get the container first, as it may be removed from the Context - // by asynchronous calls. - // This was leading to a flakey test as otherwise the container could - // be removed and end up null. + // If the container is null, then it has already completed and been removed + // from the Context by asynchronous calls. Container waitContainer = nmContext.getContainers().get(containerId); - - while ((interval-- > 0) - && !waitContainer.cloneAndGetContainerStatus() - .getState().equals(ContainerState.COMPLETE)) { - try { - LOG.info("Waiting for " + containerId + " to complete."); - Thread.sleep(1000); - } catch (InterruptedException e) { + if (waitContainer != null) { + while ((interval-- > 0) + && !waitContainer.cloneAndGetContainerStatus() + .getState().equals(ContainerState.COMPLETE)) { + try { + LOG.info("Waiting for " + containerId + " to complete."); + Thread.sleep(1000); + } catch (InterruptedException e) { + } } } // Normally, Containers will be removed from NM context after they are