From ae71f77787596d58b257cee3b822d1c11cc2b86d Mon Sep 17 00:00:00 2001
From: zhangjunfan <junfan.zhang@outlook.com>
Date: Sat, 12 Mar 2022 23:18:45 +0800
Subject: [PATCH] Fix the instability of testTonyAMStartupTimeoutShouldFail

Signed-off-by: zhangjunfan <junfan.zhang@outlook.com>
---
 .../java/com/linkedin/tony/TonyClient.java    | 10 +++++
 .../java/com/linkedin/tony/TestTonyE2E.java   | 39 +------------------
 2 files changed, 12 insertions(+), 37 deletions(-)

diff --git a/tony-core/src/main/java/com/linkedin/tony/TonyClient.java b/tony-core/src/main/java/com/linkedin/tony/TonyClient.java
index 56203ee7..ea061f90 100644
--- a/tony-core/src/main/java/com/linkedin/tony/TonyClient.java
+++ b/tony-core/src/main/java/com/linkedin/tony/TonyClient.java
@@ -168,6 +168,9 @@ public class TonyClient implements AutoCloseable {
   // For access from CLI.
   private Set<TaskInfo> taskInfos = new HashSet<>();
 
+  // Just for test cases
+  private boolean makeAppStateAlwaysAccepted = false;
+
   /**
    * Gets default hadoop application classpath from yarnConf.
    */
@@ -1070,6 +1073,10 @@ public boolean monitorApplication() throws YarnException, IOException, Interrupt
       ApplicationReport report = yarnClient.getApplicationReport(appId);
 
       YarnApplicationState appState = report.getYarnApplicationState();
+      // Just for test the config of `tony.am.startup-timeout`
+      if (makeAppStateAlwaysAccepted) {
+        appState = YarnApplicationState.ACCEPTED;
+      }
 
       FinalApplicationStatus finalApplicationStatus = report.getFinalApplicationStatus();
       initRpcClientAndLogAMUrl(report);
@@ -1476,4 +1483,7 @@ public static void main(String[] args) {
     System.exit(exitCode);
   }
 
+  public void makeAppStateAlwaysAccepted(boolean appStateAlwaysAccepted) {
+    this.makeAppStateAlwaysAccepted = appStateAlwaysAccepted;
+  }
 }
diff --git a/tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java b/tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java
index 2526c17d..073675f4 100644
--- a/tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java
+++ b/tony-core/src/test/java/com/linkedin/tony/TestTonyE2E.java
@@ -14,7 +14,6 @@
 
 import java.util.HashSet;
 
-import com.linkedin.tony.util.Utils;
 import org.apache.commons.cli.ParseException;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -34,10 +33,6 @@
 import java.net.URISyntaxException;
 import java.nio.file.Files;
 import java.util.Set;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.stream.Stream;
 
 import static com.linkedin.tony.TaskExecutor.MARK_LOST_CONNECTION_ENV_KEY;
 import static com.linkedin.tony.TonyConfigurationKeys.TASK_HEARTBEAT_INTERVAL_MS;
@@ -614,15 +609,10 @@ public void testTonyAllocationTimeoutShouldFail() throws ParseException, IOExcep
     Assert.assertEquals(exitCode, -1);
   }
 
-  @Test(timeOut = 60000)
+  @Test
   public void testTonyAMStartupTimeoutShouldFail() throws ParseException, IOException {
-    /* Creating 10 jobs submitting to Yarn to make cluster fulfill. */
-    ExecutorService executorService = Executors.newFixedThreadPool(10);
-    Stream.iterate(0, n -> n + 1)
-            .limit(10)
-            .forEach(x -> mockedTask(executorService));
-
     TonyClient client = new TonyClient(conf);
+    client.makeAppStateAlwaysAccepted(true);
     client.init(new String[]{
             "--src_dir", "tony-core/src/test/resources/scripts",
             "--executes", "python check_env_and_venv.py",
@@ -636,31 +626,6 @@ public void testTonyAMStartupTimeoutShouldFail() throws ParseException, IOExcept
     });
     int exitCode = client.start();
     Assert.assertEquals(exitCode, -1);
-
-    Utils.shutdownThreadPool(executorService);
-  }
-
-  private Future<?> mockedTask(ExecutorService service) {
-    return service.submit(() -> {
-      TonyClient client = new TonyClient(conf);
-      try {
-        client.init(new String[]{
-                "--src_dir", "tony-core/src/test/resources/scripts",
-                "--hdfs_classpath", libPath,
-                "--shell_env", "ENV_CHECK=ENV_CHECK",
-                "--container_env", Constants.SKIP_HADOOP_PATH + "=true",
-                "--python_venv", "tony-core/src/test/resources/test.zip",
-                "--conf", "tony.ps.instances=1",
-                "--conf", "tony.worker.instances=4",
-                "--conf", "tony.ps.command=python sleep_30.py",
-                "--conf", "tony.worker.command=python check_env_and_venv.py"
-        });
-        client.start();
-      } catch (Exception e) {
-        e.printStackTrace();
-      }
-      return;
-    });
   }
 
   /**