TEZ-1746. Flaky test in TestVertexImpl and TestExceptionPropagation. (zjffdu)
diff --git a/CHANGES.txt b/CHANGES.txt
index e5a4eed..359f43c 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -15,6 +15,7 @@
 ALL CHANGES:
   TEZ-1745. TestATSHistoryLoggingService::testATSHistoryLoggingServiceShutdown can be flaky.
   TEZ-1747. Increase test timeout for TestSecureShuffle.
+  TEZ-1746. Flaky test in TestVertexImpl and TestExceptionPropagation.
 
 Release 0.5.2: Unreleased
 
diff --git a/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java b/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
index 97c0693..9500c97 100644
--- a/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
+++ b/tez-dag/src/test/java/org/apache/tez/dag/app/dag/impl/TestVertexImpl.java
@@ -513,19 +513,6 @@
                         .setTaskModule("x2.y2")
                         .build()
                 )
-                .addInputs(
-                    RootInputLeafOutputProto.newBuilder()
-                        .setControllerDescriptor(
-                            TezEntityDescriptorProto.newBuilder().setClassName(
-                                "IrrelevantInitializerClassName"))
-                        .setName("input1")
-                        .setIODescriptor(
-                            TezEntityDescriptorProto.newBuilder()
-                                .setClassName("InputClazz")
-                                .build()
-                        )
-                        .build()
-                )
                 .addInEdgeId("e1")
                 .build()
         )
@@ -5072,7 +5059,7 @@
   }
   
   @Test(timeout = 5000)
-  public void testExceptionFromII_Initialize() throws AMUserCodeException {
+  public void testExceptionFromII_Initialize() throws AMUserCodeException, InterruptedException {
     useCustomInitializer = true;
     customInitializer = new EventHandlingRootInputInitializer(null, IIExceptionLocation.Initialize);
     EventHandlingRootInputInitializer initializer =
@@ -5084,6 +5071,11 @@
     VertexImplWithRunningInputInitializer v1 =
         (VertexImplWithRunningInputInitializer) vertices.get("vertex1");
     initVertex(v1);
+    // Wait for the initializer to been invoked, and fail the vertex finally.
+    while (v1.getState() != VertexState.FAILED) {
+      Thread.sleep(10);
+    }
+
     String diagnostics = StringUtils.join(v1.getDiagnostics(), ",");
     assertTrue(diagnostics.contains(IIExceptionLocation.Initialize.name()));
     Assert.assertEquals(VertexState.FAILED, v1.getState());
@@ -5206,6 +5198,10 @@
         (VertexImplWithRunningInputInitializer) vertices.get("vertex2");
 
     initVertex(v1);
+    // Wait for the initializer to be invoked - which may be a separate thread.
+    while (!initializer.initStarted.get()) {
+      Thread.sleep(10);
+    }
     startVertex(v1);  // v2 would get the state update from v1
     dispatcher.await();
     Assert.assertEquals(VertexState.RUNNING, v1.getState());
@@ -5232,6 +5228,10 @@
         (VertexImplWithRunningInputInitializer) vertices.get("vertex2");
 
     initVertex(v1);
+    // Wait for the initializer to be invoked - which may be a separate thread.
+    while (!initializer.initStarted.get()) {
+      Thread.sleep(10);
+    }
     startVertex(v1);  // v2 would get the state update from v1
     // it should be OK receive INIT_SUCCEEDED event after INIT_FAILED event
     dispatcher.getEventHandler().handle(new VertexEventRootInputInitialized(
@@ -5441,11 +5441,11 @@
 
     @Override
     public List<Event> initialize() throws Exception {
+      context.registerForVertexStateUpdates("vertex1", null);
+      initStarted.set(true);
       if (exLocation == IIExceptionLocation.Initialize) {
         throw new Exception(exLocation.name());
       }
-      context.registerForVertexStateUpdates("vertex1", null);
-      initStarted.set(true);
       lock.lock();
       try {
         if (!eventReceived.get()) {
diff --git a/tez-tests/src/test/java/org/apache/tez/test/TestExceptionPropagation.java b/tez-tests/src/test/java/org/apache/tez/test/TestExceptionPropagation.java
index 0175d7b..efb6eb6 100644
--- a/tez-tests/src/test/java/org/apache/tez/test/TestExceptionPropagation.java
+++ b/tez-tests/src/test/java/org/apache/tez/test/TestExceptionPropagation.java
@@ -212,7 +212,7 @@
    * @throws Exception
    * 
    */
-  @Test(timeout = 180000)
+  @Test(timeout = 300000)
   public void testExceptionPropagationSession() throws Exception {
     try {
       startSessionClient();