From 710cbc9bd649123cb0f742e4a91a6a216cb1ac76 Mon Sep 17 00:00:00 2001
From: Eric Yang <eyang@apache.org>
Date: Mon, 25 Mar 2019 16:34:04 -0400
Subject: [PATCH] YARN-9405.  Fixed flaky tests in TestYarnNativeServices.     
        Contributed by Prabhu Joseph

---
 .../hadoop/yarn/service/ServiceTestUtils.java |  6 +++++
 .../yarn/service/TestYarnNativeServices.java  | 22 +++++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/ServiceTestUtils.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/ServiceTestUtils.java
index b3ba58d084..a37ec759e0 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/ServiceTestUtils.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/ServiceTestUtils.java
@@ -535,6 +535,12 @@ protected void waitForServiceToBeStarted(ServiceClient client,
     waitForServiceToBeInState(client, exampleApp, ServiceState.STARTED);
   }
 
+  protected void waitForServiceToBeExpressUpgrading(ServiceClient client,
+      Service exampleApp) throws TimeoutException, InterruptedException {
+    waitForServiceToBeInState(client, exampleApp,
+        ServiceState.EXPRESS_UPGRADING);
+  }
+
   protected void waitForServiceToBeInState(ServiceClient client,
       Service exampleApp, ServiceState desiredState) throws TimeoutException,
       InterruptedException {
diff --git a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/TestYarnNativeServices.java b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/TestYarnNativeServices.java
index a22ada443c..6c38511ed7 100644
--- a/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/TestYarnNativeServices.java
+++ b/hadoop-yarn-project/hadoop-yarn/hadoop-yarn-applications/hadoop-yarn-services/hadoop-yarn-services-core/src/test/java/org/apache/hadoop/yarn/service/TestYarnNativeServices.java
@@ -439,6 +439,8 @@ public void testExpressUpgrade() throws Exception {
     component2.getConfiguration().getEnv().put("key2", "val2");
     client.actionUpgradeExpress(service);
 
+    waitForServiceToBeExpressUpgrading(client, service);
+
     // wait for upgrade to complete
     waitForServiceToBeStable(client, service);
     Service active = client.getStatus(service.getName());
@@ -859,16 +861,32 @@ private Map<String, Long> flexComponents(ServiceClient client,
   private void checkCompInstancesInOrder(ServiceClient client,
       Service exampleApp) throws IOException, YarnException,
       TimeoutException, InterruptedException {
+    waitForContainers(client, exampleApp);
     Service service = client.getStatus(exampleApp.getName());
     for (Component comp : service.getComponents()) {
       checkEachCompInstancesInOrder(comp, exampleApp.getName());
     }
   }
 
+  private void waitForContainers(ServiceClient client, Service exampleApp)
+      throws TimeoutException, InterruptedException {
+    GenericTestUtils.waitFor(() -> {
+      try {
+        Service service = client.getStatus(exampleApp.getName());
+        for (Component comp : service.getComponents()) {
+          if (comp.getContainers().size() != comp.getNumberOfContainers()) {
+            return false;
+          }
+        }
+        return true;
+      } catch (Exception e) {
+        return false;
+      }
+    }, 2000, 200000);
+  }
+
   private void checkEachCompInstancesInOrder(Component component, String
       serviceName) throws TimeoutException, InterruptedException {
-    long expectedNumInstances = component.getNumberOfContainers();
-    Assert.assertEquals(expectedNumInstances, component.getContainers().size());
     TreeSet<String> instances = new TreeSet<>();
     for (Container container : component.getContainers()) {
       instances.add(container.getComponentInstanceName());