Unmute testBanOnlyNodesWithOutstandingDescendantTasks

Author: dnhatn

server/src/internalClusterTest/java/org/elasticsearch/action/admin/cluster/node/tasks/CancellableTasksIT.java

@@ -59,6 +59,7 @@
 import org.elasticsearch.transport.TransportException;
 import org.elasticsearch.transport.TransportResponseHandler;
 import org.elasticsearch.transport.TransportService;
+import org.junit.After;
 import org.junit.Before;
 
 import java.io.IOException;
@@ -99,6 +100,16 @@ public void resetTestStates() {
         completedLatches.clear();
     }
 
+    @After
+    public void ensureAllBansRemoved() throws Exception {
+        assertBusy(() -> {
+            for (String node : internalCluster().getNodeNames()) {
+                TaskManager taskManager = internalCluster().getInstance(TransportService.class, node).getTaskManager();
+                assertThat("node " + node, taskManager.getBannedTaskIds(), empty());
+            }
+        }, 30, TimeUnit.SECONDS);
+    }
+
     static TestRequest generateTestRequest(Set<DiscoveryNode> nodes, int level, int maxLevel) {
         List<TestRequest> subRequests = new ArrayList<>();
         int lower = level == 0 ? 1 : 0;
@@ -165,21 +176,12 @@ static void allowEntireRequest(TestRequest request) {
         }
     }
 
-    void ensureAllBansRemoved() throws Exception {
-        assertBusy(() -> {
-            for (String node : internalCluster().getNodeNames()) {
-                TaskManager taskManager = internalCluster().getInstance(TransportService.class, node).getTaskManager();
-                assertThat("node " + node, taskManager.getBannedTaskIds(), empty());
-            }
-        }, 30, TimeUnit.SECONDS);
-    }
-
     public void testBanOnlyNodesWithOutstandingDescendantTasks() throws Exception {
         if (randomBoolean()) {
-            internalCluster().startNodes(randomIntBetween(1, 3));
+            internalCluster().startNodes(1);
         }
         Set<DiscoveryNode> nodes = StreamSupport.stream(clusterService().state().nodes().spliterator(), false).collect(Collectors.toSet());
-        final TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 4));
+        final TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 3));
         ActionFuture<TestResponse> rootTaskFuture = client().execute(TransportTestAction.ACTION, rootRequest);
         Set<TestRequest> pendingRequests = allowPartialRequest(rootRequest);
         TaskId rootTaskId = getRootTaskId(rootRequest);
@@ -192,28 +194,31 @@ public void testBanOnlyNodesWithOutstandingDescendantTasks() throws Exception {
                 client().admin().cluster().prepareCancelTasks().setTaskId(subTask.getTaskId()).waitForCompletion(false).get();
             }
         }
-        assertBusy(() -> {
-            for (DiscoveryNode node : nodes) {
-                TaskManager taskManager = internalCluster().getInstance(TransportService.class, node.getName()).getTaskManager();
-                Set<TaskId> expectedBans = new HashSet<>();
-                for (TestRequest req : pendingRequests) {
-                    if (req.node.equals(node)) {
-                        List<Task> childTasks = taskManager.getTasks().values().stream()
-                            .filter(t -> t.getParentTaskId() != null && t.getDescription().equals(req.taskDescription()))
-                            .collect(Collectors.toList());
-                        assertThat(childTasks, hasSize(1));
-                        CancellableTask childTask = (CancellableTask) childTasks.get(0);
-                        assertTrue(childTask.isCancelled());
-                        expectedBans.add(childTask.getParentTaskId());
+        try {
+            assertBusy(() -> {
+                for (DiscoveryNode node : nodes) {
+                    TaskManager taskManager = internalCluster().getInstance(TransportService.class, node.getName()).getTaskManager();
+                    Set<TaskId> expectedBans = new HashSet<>();
+                    for (TestRequest req : pendingRequests) {
+                        if (req.node.equals(node)) {
+                            List<Task> childTasks = taskManager.getTasks().values().stream()
+                                .filter(t -> t.getParentTaskId() != null && t.getDescription().equals(req.taskDescription()))
+                                .collect(Collectors.toList());
+                            assertThat(childTasks, hasSize(1));
+                            CancellableTask childTask = (CancellableTask) childTasks.get(0);
+                            assertTrue(childTask.isCancelled());
+                            expectedBans.add(childTask.getParentTaskId());
+                        }
                     }
+                    assertThat(taskManager.getBannedTaskIds(), equalTo(expectedBans));
                 }
-                assertThat(taskManager.getBannedTaskIds(), equalTo(expectedBans));
-            }
-        }, 30, TimeUnit.SECONDS);
-        allowEntireRequest(rootRequest);
-        cancelFuture.actionGet();
-        waitForRootTask(rootTaskFuture);
-        ensureAllBansRemoved();
+            }, 30, TimeUnit.SECONDS);
+        } finally {
+            allowEntireRequest(rootRequest);
+            cancelFuture.actionGet();
+            waitForRootTask(rootTaskFuture);
+            ensureAllBansRemoved();
+        }
     }
 
     public void testCancelTaskMultipleTimes() throws Exception {
@@ -310,7 +315,7 @@ public void testCancelOrphanedTasks() throws Exception {
 
     public void testRemoveBanParentsOnDisconnect() throws Exception {
         Set<DiscoveryNode> nodes = StreamSupport.stream(clusterService().state().nodes().spliterator(), false).collect(Collectors.toSet());
-        final TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 4));
+        final TestRequest rootRequest = generateTestRequest(nodes, 0, between(1, 3));
         client().execute(TransportTestAction.ACTION, rootRequest);
         Set<TestRequest> pendingRequests = allowPartialRequest(rootRequest);
         TaskId rootTaskId = getRootTaskId(rootRequest);