From 0700b24dd0b9fc1aeb7670e3b8224ba45107cd9f Mon Sep 17 00:00:00 2001
From: Dianna Hohensee <dianna.hohensee@elastic.co>
Date: Mon, 28 Apr 2025 14:10:28 -0400
Subject: [PATCH] Create some general test utilities (#127407)

Moving around and adding some test utilities.
---
 .../http/SearchRestCancellationIT.java          | 17 ++---------------
 .../org/elasticsearch/test/ESIntegTestCase.java | 14 ++++++++++++++
 .../elasticsearch/test/InternalTestCluster.java |  5 +++++
 3 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/qa/smoke-test-http/src/internalClusterTest/java/org/elasticsearch/http/SearchRestCancellationIT.java b/qa/smoke-test-http/src/internalClusterTest/java/org/elasticsearch/http/SearchRestCancellationIT.java
index a033582d5def..6ed1c9a1575e 100644
--- a/qa/smoke-test-http/src/internalClusterTest/java/org/elasticsearch/http/SearchRestCancellationIT.java
+++ b/qa/smoke-test-http/src/internalClusterTest/java/org/elasticsearch/http/SearchRestCancellationIT.java
@@ -12,8 +12,6 @@ import org.apache.http.entity.ContentType;
 import org.apache.http.nio.entity.NByteArrayEntity;
 import org.apache.logging.log4j.LogManager;
 import org.apache.lucene.util.SetOnce;
-import org.elasticsearch.action.admin.cluster.node.info.NodeInfo;
-import org.elasticsearch.action.admin.cluster.node.info.NodesInfoResponse;
 import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksResponse;
 import org.elasticsearch.action.bulk.BulkRequestBuilder;
 import org.elasticsearch.action.search.MultiSearchRequest;
@@ -46,7 +44,6 @@ import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CancellationException;
@@ -94,7 +91,7 @@ public class SearchRestCancellationIT extends HttpSmokeTestCase {
     }
 
     void verifyCancellationDuringQueryPhase(String searchAction, Request searchRequest) throws Exception {
-        Map<String, String> nodeIdToName = readNodesInfo();
+        Map<String, String> nodeIdToName = nodeIdsToNames();
 
         List<ScriptedBlockPlugin> plugins = initBlockFactory();
         indexTestData();
@@ -137,7 +134,7 @@ public class SearchRestCancellationIT extends HttpSmokeTestCase {
     }
 
     void verifyCancellationDuringFetchPhase(String searchAction, Request searchRequest) throws Exception {
-        Map<String, String> nodeIdToName = readNodesInfo();
+        Map<String, String> nodeIdToName = nodeIdsToNames();
 
         List<ScriptedBlockPlugin> plugins = initBlockFactory();
         indexTestData();
@@ -153,16 +150,6 @@ public class SearchRestCancellationIT extends HttpSmokeTestCase {
         expectThrows(CancellationException.class, future::actionGet);
     }
 
-    private static Map<String, String> readNodesInfo() {
-        Map<String, String> nodeIdToName = new HashMap<>();
-        NodesInfoResponse nodesInfoResponse = clusterAdmin().prepareNodesInfo().get();
-        assertFalse(nodesInfoResponse.hasFailures());
-        for (NodeInfo node : nodesInfoResponse.getNodes()) {
-            nodeIdToName.put(node.getNode().getId(), node.getNode().getName());
-        }
-        return nodeIdToName;
-    }
-
     private static void ensureSearchTaskIsCancelled(String transportAction, Function<String, String> nodeIdToName) throws Exception {
         SetOnce<TaskInfo> searchTask = new SetOnce<>();
         ListTasksResponse listTasksResponse = clusterAdmin().prepareListTasks().get();
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java
index c4bc5adfce72..a991fc78a759 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/ESIntegTestCase.java
@@ -193,6 +193,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.IdentityHashMap;
 import java.util.List;
@@ -1207,6 +1208,19 @@ public abstract class ESIntegTestCase extends ESTestCase {
         return internalCluster().getInstance(ClusterService.class, nodeName).localNode().getId();
     }
 
+    /**
+     * @return A map of the cluster node Ids to their node names.
+     */
+    public static Map<String, String> nodeIdsToNames() {
+        var names = internalCluster().getNodeNames();
+        Map<String, String> nodeIdsToNames = new HashMap<>();
+        for (var name : names) {
+            nodeIdsToNames.put(getNodeId(name), name);
+        }
+        return nodeIdsToNames;
+
+    }
+
     /**
      * Waits until at least a give number of document is visible for searchers
      *
diff --git a/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java b/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java
index e00ae7abaf0e..ab827f00f1ab 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/InternalTestCluster.java
@@ -2314,6 +2314,11 @@ public final class InternalTestCluster extends TestCluster {
         return filterNodes(nodes, NodeAndClient::isMasterEligible).size();
     }
 
+    public Set<String> masterEligibleNodeNames() {
+        var masterEligibleNodes = filterNodes(nodes, NodeAndClient::isMasterEligible);
+        return masterEligibleNodes.stream().map(nodeAndClient -> nodeAndClient.name).collect(Collectors.toSet());
+    }
+
     public void setDisruptionScheme(ServiceDisruptionScheme scheme) {
         assert activeDisruptionScheme == null
             : "there is already and active disruption [" + activeDisruptionScheme + "]. call clearDisruptionScheme first";