From 1abe6a34018eff4f0c45fd06b297da55f6f9dcd0 Mon Sep 17 00:00:00 2001
From: David Kyle <david.kyle@elastic.co>
Date: Mon, 11 Nov 2024 14:25:25 +0000
Subject: [PATCH] Delete all ingest processor pipelines

---
 .../integration/MlRestTestStateCleaner.java   | 42 +++++++++----------
 1 file changed, 21 insertions(+), 21 deletions(-)

diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
index f85788dfb9e7a..25d9509ecdc7a 100644
--- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
+++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/integration/MlRestTestStateCleaner.java
@@ -10,14 +10,15 @@
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.RestClient;
-import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.test.rest.ESRestTestCase;
 
 import java.io.IOException;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
-import java.util.stream.Collectors;
+
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.hasSize;
 
 public class MlRestTestStateCleaner {
 
@@ -30,22 +31,29 @@ public MlRestTestStateCleaner(Logger logger, RestClient adminClient) {
     }
 
     public void resetFeatures() throws IOException {
+        deletePipelinesWithInferenceProcessors();
         // This resets all features, not just ML, but they should have been getting reset between tests anyway so it shouldn't matter
         adminClient.performRequest(new Request("POST", "/_features/_reset"));
     }
 
     @SuppressWarnings("unchecked")
-    private void deleteAllTrainedModelIngestPipelines() throws IOException {
-        final Request getAllTrainedModelStats = new Request("GET", "/_ml/trained_models/_stats");
-        getAllTrainedModelStats.addParameter("size", "10000");
-        final Response trainedModelsStatsResponse = adminClient.performRequest(getAllTrainedModelStats);
+    private void deletePipelinesWithInferenceProcessors() throws IOException {
+        final Response pipelinesResponse = adminClient.performRequest(new Request("GET", "/_ingest/pipeline"));
+        final Map<String, Object> pipelines = ESRestTestCase.entityAsMap(pipelinesResponse);
+
+        var pipelinesWithInferenceProcessors = new HashSet<String>();
+        for (var entry : pipelines.entrySet()) {
+            var pipelineDef = (Map<String, Object>) entry.getValue(); // each top level object is a separate pipeline
+            var processors = (List<Map<String, Object>>) pipelineDef.get("processors");
+            for (var processor : processors) {
+                assertThat(processor.entrySet(), hasSize(1));
+                if ("inference".equals(processor.keySet().iterator().next())) {
+                    pipelinesWithInferenceProcessors.add(entry.getKey());
+                }
+            }
+        }
 
-        final List<Map<String, Object>> pipelines = (List<Map<String, Object>>) XContentMapValues.extractValue(
-            "trained_model_stats.ingest.pipelines",
-            ESRestTestCase.entityAsMap(trainedModelsStatsResponse)
-        );
-        Set<String> pipelineIds = pipelines.stream().flatMap(m -> m.keySet().stream()).collect(Collectors.toSet());
-        for (String pipelineId : pipelineIds) {
+        for (String pipelineId : pipelinesWithInferenceProcessors) {
             try {
                 adminClient.performRequest(new Request("DELETE", "/_ingest/pipeline/" + pipelineId));
             } catch (Exception ex) {
@@ -53,12 +61,4 @@ private void deleteAllTrainedModelIngestPipelines() throws IOException {
             }
         }
     }
-
-    private void waitForMlStatsIndexToInitialize() throws IOException {
-        ESRestTestCase.ensureHealth(adminClient, ".ml-stats-*", (request) -> {
-            request.addParameter("wait_for_no_initializing_shards", "true");
-            request.addParameter("level", "shards");
-            request.addParameter("timeout", "30s");
-        });
-    }
 }