HLRC: Clear ML data after client tests (#33023)

This commit duplicates the `MlRestTestStateCleaner` to make sure
all ML data is removed after each test. After implementing
the job and datafeed APIs in the HLRC, we shall replace this
implementation with one using the HLRC itself.

Closes #32993

Author: dimitris-athanasiou

client/rest-high-level/src/test/java/org/elasticsearch/client/MachineLearningIT.java

@@ -19,7 +19,6 @@
 package org.elasticsearch.client;
 
 import com.carrotsearch.randomizedtesting.generators.CodepointSetGenerator;
-import org.apache.lucene.util.LuceneTestCase.AwaitsFix;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.protocol.xpack.ml.CloseJobRequest;
 import org.elasticsearch.protocol.xpack.ml.CloseJobResponse;
@@ -33,15 +32,21 @@
 import org.elasticsearch.protocol.xpack.ml.job.config.DataDescription;
 import org.elasticsearch.protocol.xpack.ml.job.config.Detector;
 import org.elasticsearch.protocol.xpack.ml.job.config.Job;
+import org.junit.After;
 
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 
 import static org.hamcrest.Matchers.is;
 
-@AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/32993")
 public class MachineLearningIT extends ESRestHighLevelClientTestCase {
 
+    @After
+    public void cleanUp() throws IOException {
+        new MlRestTestStateCleaner(logger, client()).clearMlMetadata();
+    }
+
     public void testPutJob() throws Exception {
         String jobId = randomValidJobId();
         Job job = buildJob(jobId);

client/rest-high-level/src/test/java/org/elasticsearch/client/MlRestTestStateCleaner.java

@@ -0,0 +1,109 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.client;
+
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.common.xcontent.support.XContentMapValues;
+import org.elasticsearch.test.rest.ESRestTestCase;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * This is temporarily duplicated from the server side.
+ * @TODO Replace with an implementation using the HLRC once
+ * the APIs for managing datafeeds are implemented.
+ */
+public class MlRestTestStateCleaner {
+
+    private final Logger logger;
+    private final RestClient adminClient;
+
+    public MlRestTestStateCleaner(Logger logger, RestClient adminClient) {
+        this.logger = logger;
+        this.adminClient = adminClient;
+    }
+
+    public void clearMlMetadata() throws IOException {
+        deleteAllDatafeeds();
+        deleteAllJobs();
+        // indices will be deleted by the ESRestTestCase class
+    }
+
+    @SuppressWarnings("unchecked")
+    private void deleteAllDatafeeds() throws IOException {
+        final Request datafeedsRequest = new Request("GET", "/_xpack/ml/datafeeds");
+        datafeedsRequest.addParameter("filter_path", "datafeeds");
+        final Response datafeedsResponse = adminClient.performRequest(datafeedsRequest);
+        final List<Map<String, Object>> datafeeds =
+                (List<Map<String, Object>>) XContentMapValues.extractValue("datafeeds", ESRestTestCase.entityAsMap(datafeedsResponse));
+        if (datafeeds == null) {
+            return;
+        }
+
+        try {
+            adminClient.performRequest(new Request("POST", "/_xpack/ml/datafeeds/_all/_stop"));
+        } catch (Exception e1) {
+            logger.warn("failed to stop all datafeeds. Forcing stop", e1);
+            try {
+                adminClient.performRequest(new Request("POST", "/_xpack/ml/datafeeds/_all/_stop?force=true"));
+            } catch (Exception e2) {
+                logger.warn("Force-closing all data feeds failed", e2);
+            }
+            throw new RuntimeException(
+                    "Had to resort to force-stopping datafeeds, something went wrong?", e1);
+        }
+
+        for (Map<String, Object> datafeed : datafeeds) {
+            String datafeedId = (String) datafeed.get("datafeed_id");
+            adminClient.performRequest(new Request("DELETE", "/_xpack/ml/datafeeds/" + datafeedId));
+        }
+    }
+
+    private void deleteAllJobs() throws IOException {
+        final Request jobsRequest = new Request("GET", "/_xpack/ml/anomaly_detectors");
+        jobsRequest.addParameter("filter_path", "jobs");
+        final Response response = adminClient.performRequest(jobsRequest);
+        @SuppressWarnings("unchecked")
+        final List<Map<String, Object>> jobConfigs =
+                (List<Map<String, Object>>) XContentMapValues.extractValue("jobs", ESRestTestCase.entityAsMap(response));
+        if (jobConfigs == null) {
+            return;
+        }
+
+        try {
+            adminClient.performRequest(new Request("POST", "/_xpack/ml/anomaly_detectors/_all/_close"));
+        } catch (Exception e1) {
+            logger.warn("failed to close all jobs. Forcing closed", e1);
+            try {
+                adminClient.performRequest(new Request("POST", "/_xpack/ml/anomaly_detectors/_all/_close?force=true"));
+            } catch (Exception e2) {
+                logger.warn("Force-closing all jobs failed", e2);
+            }
+            throw new RuntimeException("Had to resort to force-closing jobs, something went wrong?",
+                    e1);
+        }
+
+        for (Map<String, Object> jobConfig : jobConfigs) {
+            String jobId = (String) jobConfig.get("job_id");
+            adminClient.performRequest(new Request("DELETE", "/_xpack/ml/anomaly_detectors/" + jobId));
+        }
+    }
+}

client/rest-high-level/src/test/java/org/elasticsearch/client/documentation/MlClientDocumentationIT.java

@@ -22,6 +22,7 @@
 import org.elasticsearch.action.LatchedActionListener;
 import org.elasticsearch.client.ESRestHighLevelClientTestCase;
 import org.elasticsearch.client.MachineLearningIT;
+import org.elasticsearch.client.MlRestTestStateCleaner;
 import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.RestHighLevelClient;
 import org.elasticsearch.common.unit.TimeValue;
@@ -37,7 +38,9 @@
 import org.elasticsearch.protocol.xpack.ml.job.config.DataDescription;
 import org.elasticsearch.protocol.xpack.ml.job.config.Detector;
 import org.elasticsearch.protocol.xpack.ml.job.config.Job;
+import org.junit.After;
 
+import java.io.IOException;
 import java.util.Collections;
 import java.util.Date;
 import java.util.List;
@@ -48,6 +51,11 @@
 
 public class MlClientDocumentationIT extends ESRestHighLevelClientTestCase {
 
+    @After
+    public void cleanUp() throws IOException {
+        new MlRestTestStateCleaner(logger, client()).clearMlMetadata();
+    }
+
     public void testCreateJob() throws Exception {
         RestHighLevelClient client = highLevelClient();