From 13d3b353c6c5fb6047d7241efc653863417e6733 Mon Sep 17 00:00:00 2001
From: Dimitris Athanasiou <dimitris@elastic.co>
Date: Thu, 20 Apr 2017 14:03:38 +0100
Subject: [PATCH] [TEST] Wait for pending tasks on test clean up...
 (elastic/x-pack-elasticsearch#1137)

.. in MlNativeAutodetectIntegTestCase.

We wait for pending tasks to finish from REST tests but not
from the ones using a native autodetect. This commit adds the
waiting in those tests too.

Relates elastic/x-pack-elasticsearch#1136

Original commit: elastic/x-pack-elasticsearch@a7a5455c7883f4966adddaea00447a2446f11ae7
---
 .../MlNativeAutodetectIntegTestCase.java      | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeAutodetectIntegTestCase.java b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeAutodetectIntegTestCase.java
index 527459fd5d4..9250a29ed29 100644
--- a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeAutodetectIntegTestCase.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlNativeAutodetectIntegTestCase.java
@@ -5,9 +5,11 @@
  */
 package org.elasticsearch.xpack.ml.integration;
 
+import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksRequest;
 import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.network.NetworkModule;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.test.SecurityIntegTestCase;
 import org.elasticsearch.xpack.XPackSettings;
@@ -63,6 +65,12 @@ abstract class MlNativeAutodetectIntegTestCase extends SecurityIntegTestCase {
     }
 
     protected void cleanUp() {
+        cleanUpDatafeeds();
+        cleanUpJobs();
+        waitForPendingTasks();
+    }
+
+    private void cleanUpDatafeeds() {
         for (DatafeedConfig datafeed : datafeeds) {
             try {
                 stopDatafeed(datafeed.getId());
@@ -75,6 +83,9 @@ abstract class MlNativeAutodetectIntegTestCase extends SecurityIntegTestCase {
                 // ignore
             }
         }
+    }
+
+    private void cleanUpJobs() {
         for (Job.Builder job : jobs) {
             try {
                 closeJob(job.getId());
@@ -89,6 +100,18 @@ abstract class MlNativeAutodetectIntegTestCase extends SecurityIntegTestCase {
         }
     }
 
+    private void waitForPendingTasks() {
+        ListTasksRequest listTasksRequest = new ListTasksRequest();
+        listTasksRequest.setWaitForCompletion(true);
+        listTasksRequest.setDetailed(true);
+        listTasksRequest.setTimeout(TimeValue.timeValueSeconds(10));
+        try {
+            admin().cluster().listTasks(listTasksRequest).get();
+        } catch (Exception e) {
+            throw new AssertionError("Failed to wait for pending tasks to complete", e);
+        }
+    }
+
     protected void registerJob(Job.Builder job) {
         if (jobs.add(job) == false) {
             throw new IllegalArgumentException("job [" + job.getId() + "] is already registered");