[8.x] Improve CrossClusterAsyncEnrichStopIT test (#122432)

muted-tests.yml

@@ -531,9 +531,6 @@ tests:
   issue: https://github.com/elastic/elasticsearch/issues/121483
 - class: org.elasticsearch.xpack.application.FullClusterRestartIT
   issue: https://github.com/elastic/elasticsearch/issues/121935
-- class: org.elasticsearch.xpack.esql.action.CrossClusterAsyncEnrichStopIT
-  method: testEnrichAfterStop
-  issue: https://github.com/elastic/elasticsearch/issues/121994
 - class: org.elasticsearch.xpack.esql.qa.mixed.MixedClusterEsqlSpecIT
   method: test {date_nanos.Bucket Date nanos by 10 minutes SYNC}
   issue: https://github.com/elastic/elasticsearch/issues/122273

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractCrossClusterTestCase.java

@@ -14,9 +14,11 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.compute.operator.DriverTaskRunner;
 import org.elasticsearch.compute.operator.exchange.ExchangeService;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.test.AbstractMultiClustersTestCase;
 import org.elasticsearch.test.FailingFieldPlugin;
 import org.elasticsearch.test.XContentTestUtils;
@@ -273,4 +275,8 @@ protected EsqlQueryResponse runQuery(String query, Boolean ccsMetadataInResponse
         }
         return runQuery(request);
     }
+
+    static List<TaskInfo> getDriverTasks(Client client) {
+        return client.admin().cluster().prepareListTasks().setActions(DriverTaskRunner.ACTION_NAME).setDetailed(true).get().getTasks();
+    }
 }

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClusterAsyncEnrichStopIT.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.xcontent.json.JsonXContent;
 import org.elasticsearch.xpack.core.async.AsyncStopRequest;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -28,10 +29,13 @@
 
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
+import static org.elasticsearch.xpack.esql.action.AbstractCrossClusterTestCase.getDriverTasks;
 import static org.elasticsearch.xpack.esql.action.EsqlAsyncTestUtils.deleteAsyncId;
 import static org.elasticsearch.xpack.esql.action.EsqlAsyncTestUtils.startAsyncQuery;
 import static org.elasticsearch.xpack.esql.action.EsqlAsyncTestUtils.waitForCluster;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.not;
 
 // This tests if enrich after stop works correctly
 public class CrossClusterAsyncEnrichStopIT extends AbstractEnrichBasedCrossClusterTestCase {
@@ -87,10 +91,23 @@ public void testEnrichAfterStop() throws Exception {
         // wait until c1 is done
         waitForCluster(client(), "c1", asyncExecutionId);
         waitForCluster(client(), LOCAL_CLUSTER, asyncExecutionId);
+        // wait until remote reduce task starts on c2
+        assertBusy(() -> {
+            List<TaskInfo> tasks = getDriverTasks(client(REMOTE_CLUSTER_2));
+            List<TaskInfo> reduceTasks = tasks.stream().filter(t -> t.description().contains("_LuceneSourceOperator") == false).toList();
+            assertThat(reduceTasks, not(empty()));
+        });
 
         // Run the stop request
         var stopRequest = new AsyncStopRequest(asyncExecutionId);
         var stopAction = client().execute(EsqlAsyncStopAction.INSTANCE, stopRequest);
+        // wait until remote reduce tasks are gone
+        assertBusy(() -> {
+            List<TaskInfo> tasks = getDriverTasks(client(REMOTE_CLUSTER_2));
+            List<TaskInfo> reduceTasks = tasks.stream().filter(t -> t.description().contains("_LuceneSourceOperator") == false).toList();
+            assertThat(reduceTasks, empty());
+        });
+
         // Allow the processing to proceed
         SimplePauseFieldPlugin.allowEmitting.countDown();
 

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClusterAsyncQueryStopIT.java

@@ -9,8 +9,6 @@
 
 import org.elasticsearch.Build;
 import org.elasticsearch.action.ActionFuture;
-import org.elasticsearch.client.internal.Client;
-import org.elasticsearch.compute.operator.DriverTaskRunner;
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.tasks.TaskInfo;
 import org.elasticsearch.xpack.core.async.AsyncStopRequest;
@@ -244,8 +242,4 @@ public void testStopQueryLocalNoRemotes() throws Exception {
             assertAcked(deleteAsyncId(client(), asyncExecutionId));
         }
     }
-
-    private static List<TaskInfo> getDriverTasks(Client client) {
-        return client.admin().cluster().prepareListTasks().setActions(DriverTaskRunner.ACTION_NAME).setDetailed(true).get().getTasks();
-    }
 }