From f5081a9a5d51fe781da97f07ac4dca8dd49c98e2 Mon Sep 17 00:00:00 2001
From: Ahmed Hussein <ahussein@apache.org>
Date: Mon, 4 May 2020 10:13:51 -0500
Subject: [PATCH] YARN-10256. Refactor
 TestContainerSchedulerQueuing.testContainerUpdateExecTypeGuaranteedToOpportunistic
 (Ahmed Hussein via jeagles)

Signed-off-by: Jonathan Eagles <jeagles@gmail.com>
---
 .../TestContainerSchedulerQueuing.java        | 50 +++++++------------
 1 file changed, 18 insertions(+), 32 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/scheduler/TestContainerSchedulerQueuing.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/scheduler/TestContainerSchedulerQueuing.java
index 530a8e7260..b21850cbcf 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/scheduler/TestContainerSchedulerQueuing.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-server/hadoop-yarn-server-nodemanager/src/test/java/org/apache/hadoop/yarn/server/nodemanager/containermanager/scheduler/TestContainerSchedulerQueuing.java
@@ -28,7 +28,6 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
-import com.google.common.base.Supplier;
 import org.apache.hadoop.fs.UnsupportedFileSystemException;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.test.GenericTestUtils;
@@ -1291,37 +1290,24 @@ public void testContainerUpdateExecTypeGuaranteedToOpportunistic()
         1, updateResponse.getSuccessfullyUpdatedContainers().size());
     Assert.assertTrue(updateResponse.getFailedRequests().isEmpty());
 
-    GetContainerStatusesRequest statRequest =
-        GetContainerStatusesRequest.newInstance(Collections.singletonList(cId));
-    GenericTestUtils.waitFor(
-        new Supplier<Boolean>() {
-          @Override
-          public Boolean get() {
-            try {
-              List<ContainerStatus> containerStatuses = containerManager
-                  .getContainerStatuses(statRequest).getContainerStatuses();
-              Assert.assertEquals(1, containerStatuses.size());
+    final GetContainerStatusesRequest statRequest =
+            GetContainerStatusesRequest.newInstance(
+                    Collections.singletonList(cId));
+    final org.apache.hadoop.yarn.api.records.ContainerState expectedState =
+            org.apache.hadoop.yarn.api.records.ContainerState.RUNNING;
 
-              ContainerStatus status = containerStatuses.get(0);
-              Assert.assertEquals(
-                  org.apache.hadoop.yarn.api.records.ContainerState.RUNNING,
-                  status.getState());
-
-              return status.getExecutionType() == ExecutionType.OPPORTUNISTIC;
-            } catch (Exception ex) {
-              throw new RuntimeException(ex);
-            }
-          }
-        }, 100, 10000);
-    List<ContainerStatus> containerStatuses = containerManager
-        .getContainerStatuses(statRequest).getContainerStatuses();
-    Assert.assertEquals(1, containerStatuses.size());
-    for (ContainerStatus status : containerStatuses) {
-      Assert.assertEquals(
-          org.apache.hadoop.yarn.api.records.ContainerState.RUNNING,
-          status.getState());
-      Assert
-          .assertEquals(ExecutionType.OPPORTUNISTIC, status.getExecutionType());
-    }
+    GenericTestUtils.waitFor(() -> {
+      List<ContainerStatus> containerStatuses;
+      try {
+        containerStatuses = containerManager
+                .getContainerStatuses(statRequest).getContainerStatuses();
+      } catch (YarnException | IOException e) {
+        return false;
+      }
+      Assert.assertEquals(1, containerStatuses.size());
+      ContainerStatus status = containerStatuses.get(0);
+      return (status.getState() == expectedState
+              && status.getExecutionType() == ExecutionType.OPPORTUNISTIC);
+    }, 20, 10000);
   }
 }