From b24dc2c5417af62a007ef0104d8bb6d62488a7b7 Mon Sep 17 00:00:00 2001
From: Jim Ferenczi <jim.ferenczi@elastic.co>
Date: Fri, 11 Jan 2019 09:47:46 +0100
Subject: [PATCH] [TEST] Awaits tasks termination in the RestHighLevelClient
 tests (#37302)

This change ensures that TasksIT#testGetValidTask and ReindexIT#testReindexTask
don't leave a non-completed task on the cluster when they finish.

Closes #35644
---
 .../org/elasticsearch/client/ReindexIT.java    | 18 ++++++------------
 .../java/org/elasticsearch/client/TasksIT.java |  7 +++++--
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/client/rest-high-level/src/test/java/org/elasticsearch/client/ReindexIT.java b/client/rest-high-level/src/test/java/org/elasticsearch/client/ReindexIT.java
index 73594d8a878..cfdd29cdfbf 100644
--- a/client/rest-high-level/src/test/java/org/elasticsearch/client/ReindexIT.java
+++ b/client/rest-high-level/src/test/java/org/elasticsearch/client/ReindexIT.java
@@ -23,6 +23,7 @@ import org.elasticsearch.action.bulk.BulkRequest;
 import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.action.support.WriteRequest;
 import org.elasticsearch.client.tasks.TaskSubmissionResponse;
+import org.elasticsearch.common.CheckedRunnable;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.query.IdsQueryBuilder;
@@ -32,7 +33,6 @@ import org.elasticsearch.rest.RestStatus;
 
 import java.io.IOException;
 import java.util.Collections;
-import java.util.function.BooleanSupplier;
 
 public class ReindexIT extends ESRestHighLevelClientTestCase {
 
@@ -82,7 +82,7 @@ public class ReindexIT extends ESRestHighLevelClientTestCase {
         }
     }
 
-    public void testReindexTask() throws IOException, InterruptedException {
+    public void testReindexTask() throws Exception {
         final String sourceIndex = "source123";
         final String destinationIndex = "dest2";
         {
@@ -118,20 +118,14 @@ public class ReindexIT extends ESRestHighLevelClientTestCase {
             String taskId = reindexSubmission.getTask(); // <3>
             // end::submit-reindex-task
 
-            BooleanSupplier hasUpgradeCompleted = checkCompletionStatus(taskId);
-            awaitBusy(hasUpgradeCompleted);
+            assertBusy(checkCompletionStatus(client(), taskId));
         }
     }
 
-    private BooleanSupplier checkCompletionStatus(String taskId) {
+    static CheckedRunnable<Exception> checkCompletionStatus(RestClient client, String taskId) {
         return () -> {
-            try {
-                Response response = client().performRequest(new Request("GET", "/_tasks/" + taskId));
-                return (boolean) entityAsMap(response).get("completed");
-            } catch (IOException e) {
-                fail(e.getMessage());
-                return false;
-            }
+            Response response = client.performRequest(new Request("GET", "/_tasks/" + taskId));
+            assertTrue((boolean) entityAsMap(response).get("completed"));
         };
     }
 }
diff --git a/client/rest-high-level/src/test/java/org/elasticsearch/client/TasksIT.java b/client/rest-high-level/src/test/java/org/elasticsearch/client/TasksIT.java
index 28b909df5d4..5dc168eadfe 100644
--- a/client/rest-high-level/src/test/java/org/elasticsearch/client/TasksIT.java
+++ b/client/rest-high-level/src/test/java/org/elasticsearch/client/TasksIT.java
@@ -72,7 +72,7 @@ public class TasksIT extends ESRestHighLevelClientTestCase {
         assertTrue("List tasks were not found", listTasksFound);
     }
     
-    public void testGetValidTask() throws IOException {
+    public void testGetValidTask() throws Exception {
 
         // Run a Reindex to create a task
 
@@ -112,7 +112,10 @@ public class TasksIT extends ESRestHighLevelClientTestCase {
         TaskInfo info = taskResponse.getTaskInfo();
         assertTrue(info.isCancellable());
         assertEquals("reindex from [source1] to [dest][_doc]", info.getDescription());
-        assertEquals("indices:data/write/reindex", info.getAction());                
+        assertEquals("indices:data/write/reindex", info.getAction());
+        if (taskResponse.isCompleted() == false) {
+            assertBusy(ReindexIT.checkCompletionStatus(client(), taskId.toString()));
+        }
     }    
     
     public void testGetInvalidTask() throws IOException {