From 7584fbf4cbafd34fac4b362cefe4e06cec16a2af Mon Sep 17 00:00:00 2001
From: Varun Saxena <varunsaxena@apache.org>
Date: Sat, 19 Nov 2016 01:37:59 +0530
Subject: [PATCH] MAPREDUCE-6801. Fix flaky TestKill.testKillJob (Haibo Chen
 via Varun Saxena)

---
 .../org/apache/hadoop/mapreduce/v2/app/TestKill.java     | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/test/java/org/apache/hadoop/mapreduce/v2/app/TestKill.java b/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/test/java/org/apache/hadoop/mapreduce/v2/app/TestKill.java
index aae591e33d3..07146478028 100644
--- a/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/test/java/org/apache/hadoop/mapreduce/v2/app/TestKill.java
+++ b/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-app/src/test/java/org/apache/hadoop/mapreduce/v2/app/TestKill.java
@@ -22,6 +22,7 @@
 import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 
+import org.apache.hadoop.service.Service;
 import org.junit.Assert;
 
 import org.apache.hadoop.conf.Configuration;
@@ -66,8 +67,8 @@ public void testKillJob() throws Exception {
     Job job = app.submit(new Configuration());
     
     //wait and vailidate for Job to become RUNNING
-    app.waitForState(job, JobState.RUNNING);
-    
+    app.waitForInternalState((JobImpl) job, JobStateInternal.RUNNING);
+
     //send the kill signal to Job
     app.getContext().getEventHandler().handle(
         new JobEvent(job.getID(), JobEventType.JOB_KILL));
@@ -77,6 +78,10 @@ public void testKillJob() throws Exception {
 
     //wait and validate for Job to be KILLED
     app.waitForState(job, JobState.KILLED);
+    // make sure all events are processed. The AM is stopped
+    // only when all tasks and task attempts have been killed
+    app.waitForState(Service.STATE.STOPPED);
+
     Map<TaskId,Task> tasks = job.getTasks();
     Assert.assertEquals("No of tasks is not correct", 1, 
         tasks.size());