TEZ-4475: VertexStatus is missing in TestLocalMode if DAG finishes too quickly - causing NPE in unit test

Author: abstractdog

tez-api/src/main/java/org/apache/tez/dag/api/client/DAGClientImpl.java

@@ -303,6 +303,9 @@ public VertexStatus getVertexStatus(String vertexName, Set<StatusGetOpts> status
 
   protected VertexStatus getVertexStatusInternal(Set<StatusGetOpts> statusOptions, String vertexName)
       throws IOException, TezException {
+    LOG.debug("Called getVertexStatusInternal for {}, dagCompleted: {}, vertex status in cache: {}", vertexName,
+        dagCompleted, cachedVertexStatus.containsKey(vertexName));
+
     if (!dagCompleted) {
       VertexStatus vertexStatus = getVertexStatusViaAM(vertexName, statusOptions);
 

tez-tests/src/test/java/org/apache/tez/test/TestLocalMode.java

@@ -130,7 +130,7 @@ public void testMultipleClientsWithSession() throws TezException, InterruptedExc
     DAG dag1 = createSimpleDAG("dag1", SleepProcessor.class.getName());
 
     DAGClient dagClient1 = tezClient1.submitDAG(dag1);
-    dagClient1.waitForCompletion();
+    dagClient1.waitForCompletionWithStatusUpdates(null);
     assertEquals(DAGStatus.State.SUCCEEDED, dagClient1.getDAGStatus(null).getState());
     assertEquals(VertexStatus.State.SUCCEEDED,
         dagClient1.getVertexStatus(SleepProcessor.SLEEP_VERTEX_NAME, null).getState());
@@ -143,7 +143,7 @@ public void testMultipleClientsWithSession() throws TezException, InterruptedExc
     TezClient tezClient2 = TezClient.create("commonName", tezConf2, true);
     tezClient2.start();
     DAGClient dagClient2 = tezClient2.submitDAG(dag2);
-    dagClient2.waitForCompletion();
+    dagClient2.waitForCompletionWithStatusUpdates(null);
     assertEquals(DAGStatus.State.SUCCEEDED, dagClient2.getDAGStatus(null).getState());
     assertEquals(VertexStatus.State.SUCCEEDED,
         dagClient2.getVertexStatus(SleepProcessor.SLEEP_VERTEX_NAME, null).getState());
@@ -162,7 +162,7 @@ public void testMultipleClientsWithoutSession() throws TezException, Interrupted
     DAG dag1 = createSimpleDAG("dag1", SleepProcessor.class.getName());
 
     DAGClient dagClient1 = tezClient1.submitDAG(dag1);
-    dagClient1.waitForCompletion();
+    dagClient1.waitForCompletionWithStatusUpdates(null);
     assertEquals(DAGStatus.State.SUCCEEDED, dagClient1.getDAGStatus(null).getState());
     assertEquals(VertexStatus.State.SUCCEEDED,
         dagClient1.getVertexStatus(SleepProcessor.SLEEP_VERTEX_NAME, null).getState());
@@ -175,7 +175,7 @@ public void testMultipleClientsWithoutSession() throws TezException, Interrupted
     TezClient tezClient2 = TezClient.create("commonName", tezConf2, false);
     tezClient2.start();
     DAGClient dagClient2 = tezClient2.submitDAG(dag2);
-    dagClient2.waitForCompletion();
+    dagClient2.waitForCompletionWithStatusUpdates(null);
     assertEquals(DAGStatus.State.SUCCEEDED, dagClient2.getDAGStatus(null).getState());
     assertEquals(VertexStatus.State.SUCCEEDED,
         dagClient2.getVertexStatus(SleepProcessor.SLEEP_VERTEX_NAME, null).getState());
@@ -195,7 +195,7 @@ public void testNoSysExitOnSuccessfulDAG() throws TezException, InterruptedExcep
     DAG dag1 = createSimpleDAG("dag1", SleepProcessor.class.getName());
 
     DAGClient dagClient1 = tezClient1.submitDAG(dag1);
-    dagClient1.waitForCompletion();
+    dagClient1.waitForCompletionWithStatusUpdates(null);
     assertEquals(DAGStatus.State.SUCCEEDED, dagClient1.getDAGStatus(null).getState());
     assertEquals(VertexStatus.State.SUCCEEDED,
         dagClient1.getVertexStatus(SleepProcessor.SLEEP_VERTEX_NAME, null).getState());
@@ -217,7 +217,7 @@ public void testNoSysExitOnFailinglDAG() throws TezException, InterruptedExcepti
     DAG dag1 = createSimpleDAG("dag1", FailingProcessor.class.getName());
 
     DAGClient dagClient1 = tezClient1.submitDAG(dag1);
-    dagClient1.waitForCompletion();
+    dagClient1.waitForCompletionWithStatusUpdates(null);
     assertEquals(DAGStatus.State.FAILED, dagClient1.getDAGStatus(null).getState());
     assertEquals(VertexStatus.State.FAILED,
         dagClient1.getVertexStatus(SleepProcessor.SLEEP_VERTEX_NAME, null).getState());
@@ -290,7 +290,7 @@ public void testMultiDAGsOnSession() throws IOException, TezException, Interrupt
         dagClients[i] = tezClient.submitDAG(dag);
 
         // wait to finish
-        DAGStatus dagStatus = dagClients[i].waitForCompletion();
+        DAGStatus dagStatus = dagClients[i].waitForCompletionWithStatusUpdates(null);
         if (dagStatus.getState() != DAGStatus.State.SUCCEEDED) {
           fail("Iteration " + i + " failed with diagnostics: "
               + dagStatus.getDiagnostics());