From 7cb2be2e5cee88ea159c4b14b256820923656bce Mon Sep 17 00:00:00 2001
From: Colin Goodheart-Smithe <colings86@users.noreply.github.com>
Date: Wed, 19 Apr 2017 17:36:30 +0100
Subject: [PATCH] Adds a check to wait for active tasks for XPackRestIT
 (elastic/x-pack-elasticsearch#964)

* Adds a check to wait for active tasks for XPackRestIT

* uses test logger

* Change to use assertBusy instead of awaitBusy

* fixes failures with active tasks remaining

* Moves wait for pending tasks into MlRestTestStateCleaner

* remove unecessary log line

Original commit: elastic/x-pack-elasticsearch@1f098dbb64e99a3581af0961449a36bca1e8294a
---
 .../xpack/ml/integration/MlJobIT.java         |  3 +-
 .../integration/MlRestTestStateCleaner.java   | 41 ++++++++++++++++++-
 .../xpack/test/rest/XPackRestIT.java          |  2 +-
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlJobIT.java b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlJobIT.java
index 4659bd73c76..2791245ba98 100644
--- a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlJobIT.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlJobIT.java
@@ -18,7 +18,6 @@ import org.elasticsearch.xpack.ml.job.persistence.AnomalyDetectorsIndex;
 import org.junit.After;
 
 import java.io.BufferedReader;
-import java.io.IOException;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
@@ -659,7 +658,7 @@ public class MlJobIT extends ESRestTestCase {
     }
 
     @After
-    public void clearMlState() throws IOException {
+    public void clearMlState() throws Exception {
         new MlRestTestStateCleaner(logger, adminClient(), this).clearMlMetadata();
     }
 }
diff --git a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlRestTestStateCleaner.java b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlRestTestStateCleaner.java
index f44ce2b2625..3fd8370667a 100644
--- a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlRestTestStateCleaner.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlRestTestStateCleaner.java
@@ -5,16 +5,25 @@
  */
 package org.elasticsearch.xpack.ml.integration;
 
+import org.apache.http.HttpStatus;
 import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksAction;
+import org.elasticsearch.client.Response;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
+import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.rest.ESRestTestCase;
 
+import java.io.BufferedReader;
 import java.io.IOException;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
+import static org.junit.Assert.assertEquals;
+
 /*
  * NOTE: a copy if this file resides in :x-pack-elasticsearch:qa:smoke-test-ml-with-security
  *
@@ -34,12 +43,42 @@ public class MlRestTestStateCleaner {
         this.testCase = testCase;
     }
 
-    public void clearMlMetadata() throws IOException {
+    public void clearMlMetadata() throws Exception {
         deleteAllDatafeeds();
         deleteAllJobs();
+        waitForPendingTasks();
         deleteDotML();
     }
 
+    private void waitForPendingTasks() throws Exception {
+        ESTestCase.assertBusy(() -> {
+            try {
+                Response response = adminClient.performRequest("GET", "/_cat/tasks");
+                // Check to see if there are tasks still active. We exclude the
+                // list tasks
+                // actions tasks form this otherwise we will always fail
+                if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
+                    try (BufferedReader responseReader = new BufferedReader(
+                            new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8))) {
+                        int activeTasks = 0;
+                        String line;
+                        StringBuilder tasksListString = new StringBuilder();
+                        while ((line = responseReader.readLine()) != null) {
+                            tasksListString.append(line);
+                            tasksListString.append('\n');
+                            if (line.startsWith(ListTasksAction.NAME) == false) {
+                                activeTasks++;
+                            }
+                        }
+                        assertEquals(activeTasks + " active tasks found:\n" + tasksListString, 0, activeTasks);
+                    }
+                }
+            } catch (IOException e) {
+                throw new AssertionError("Error getting active tasks list", e);
+            }
+        });
+    }
+
     @SuppressWarnings("unchecked")
     private void deleteAllDatafeeds() throws IOException {
         Map<String, Object> clusterStateAsMap = testCase.entityAsMap(adminClient.performRequest("GET", "/_cluster/state",
diff --git a/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestIT.java b/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestIT.java
index 1f2ac9efd13..e983c8d250f 100644
--- a/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestIT.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestIT.java
@@ -38,7 +38,7 @@ import static java.util.Collections.singletonMap;
 public class XPackRestIT extends XPackRestTestCase {
 
     @After
-    public void clearMlState() throws IOException {
+    public void clearMlState() throws Exception {
         new MlRestTestStateCleaner(logger, adminClient(), this).clearMlMetadata();
     }