From bda415b7fdf4a73091a198339e5f1660c1378029 Mon Sep 17 00:00:00 2001
From: Nhat Nguyen <nhat.nguyen@elastic.co>
Date: Sat, 30 Nov 2024 20:09:08 -0800
Subject: [PATCH] Fix CCS cancellation test (#117790)

We should have checked that all drivers were canceled, not cancellable
(which is always true), before unblocking the compute tasks.

Closes #117568
---
 muted-tests.yml                               |  3 --
 .../action/CrossClustersCancellationIT.java   | 29 ++++++++++---------
 2 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/muted-tests.yml b/muted-tests.yml
index d5e2dbd84cb4a..b82e95ea26890 100644
--- a/muted-tests.yml
+++ b/muted-tests.yml
@@ -234,9 +234,6 @@ tests:
 - class: org.elasticsearch.search.ccs.CrossClusterIT
   method: testCancel
   issue: https://github.com/elastic/elasticsearch/issues/108061
-- class: org.elasticsearch.xpack.esql.action.CrossClustersCancellationIT
-  method: testCancel
-  issue: https://github.com/elastic/elasticsearch/issues/117568
 
 # Examples:
 #
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersCancellationIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersCancellationIT.java
index c426e0f528eab..5ffc92636b272 100644
--- a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersCancellationIT.java
+++ b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClustersCancellationIT.java
@@ -179,19 +179,22 @@ public void testCancel() throws Exception {
         });
         var cancelRequest = new CancelTasksRequest().setTargetTaskId(rootTasks.get(0).taskId()).setReason("proxy timeout");
         client().execute(TransportCancelTasksAction.TYPE, cancelRequest);
-        assertBusy(() -> {
-            List<TaskInfo> drivers = client(REMOTE_CLUSTER).admin()
-                .cluster()
-                .prepareListTasks()
-                .setActions(DriverTaskRunner.ACTION_NAME)
-                .get()
-                .getTasks();
-            assertThat(drivers.size(), greaterThanOrEqualTo(1));
-            for (TaskInfo driver : drivers) {
-                assertTrue(driver.cancellable());
-            }
-        });
-        PauseFieldPlugin.allowEmitting.countDown();
+        try {
+            assertBusy(() -> {
+                List<TaskInfo> drivers = client(REMOTE_CLUSTER).admin()
+                    .cluster()
+                    .prepareListTasks()
+                    .setActions(DriverTaskRunner.ACTION_NAME)
+                    .get()
+                    .getTasks();
+                assertThat(drivers.size(), greaterThanOrEqualTo(1));
+                for (TaskInfo driver : drivers) {
+                    assertTrue(driver.cancelled());
+                }
+            });
+        } finally {
+            PauseFieldPlugin.allowEmitting.countDown();
+        }
         Exception error = expectThrows(Exception.class, requestFuture::actionGet);
         assertThat(error.getMessage(), containsString("proxy timeout"));
     }