[ML][HLRC] Replace REST-based ML test cleanup with the ML client (#34109)

Now that all basic APIs for managing jobs and datafeeds have been
implemented we replace the duplicated `MlRestTestStateCleaner`
with an implementation that uses the HLRC Machine Learning client
itself.

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

@@ -147,7 +147,7 @@ public class MachineLearningGetResultsIT extends ESRestHighLevelClientTestCase {
 
     @After
     public void deleteJob() throws IOException {
-        new MlRestTestStateCleaner(logger, client()).clearMlMetadata();
+        new MlTestStateCleaner(logger, highLevelClient().machineLearning()).clearMlMetadata();
     }
 
     public void testGetCategories() throws IOException {

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.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
 import org.elasticsearch.action.bulk.BulkRequest;
@@ -93,7 +92,7 @@ public class MachineLearningIT extends ESRestHighLevelClientTestCase {
 
     @After
     public void cleanUp() throws IOException {
-        new MlRestTestStateCleaner(logger, client()).clearMlMetadata();
+        new MlTestStateCleaner(logger, highLevelClient().machineLearning()).clearMlMetadata();
     }
 
     public void testPutJob() throws Exception {

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

@@ -1,109 +0,0 @@
-/*
- * 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/MlTestStateCleaner.java

@@ -0,0 +1,102 @@
+/*
+ * 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.client.ml.CloseJobRequest;
+import org.elasticsearch.client.ml.DeleteDatafeedRequest;
+import org.elasticsearch.client.ml.DeleteJobRequest;
+import org.elasticsearch.client.ml.GetDatafeedRequest;
+import org.elasticsearch.client.ml.GetDatafeedResponse;
+import org.elasticsearch.client.ml.GetJobRequest;
+import org.elasticsearch.client.ml.GetJobResponse;
+import org.elasticsearch.client.ml.StopDatafeedRequest;
+import org.elasticsearch.client.ml.datafeed.DatafeedConfig;
+import org.elasticsearch.client.ml.job.config.Job;
+
+import java.io.IOException;
+
+/**
+ * Cleans up and ML resources created during tests
+ */
+public class MlTestStateCleaner {
+
+    private final Logger logger;
+    private final MachineLearningClient mlClient;
+
+    public MlTestStateCleaner(Logger logger, MachineLearningClient mlClient) {
+        this.logger = logger;
+        this.mlClient = mlClient;
+    }
+
+    public void clearMlMetadata() throws IOException {
+        deleteAllDatafeeds();
+        deleteAllJobs();
+    }
+
+    private void deleteAllDatafeeds() throws IOException {
+        stopAllDatafeeds();
+
+        GetDatafeedResponse getDatafeedResponse = mlClient.getDatafeed(GetDatafeedRequest.getAllDatafeedsRequest(), RequestOptions.DEFAULT);
+        for (DatafeedConfig datafeed : getDatafeedResponse.datafeeds()) {
+            mlClient.deleteDatafeed(new DeleteDatafeedRequest(datafeed.getId()), RequestOptions.DEFAULT);
+        }
+    }
+
+    private void stopAllDatafeeds() {
+        StopDatafeedRequest stopAllDatafeedsRequest = StopDatafeedRequest.stopAllDatafeedsRequest();
+        try {
+            mlClient.stopDatafeed(stopAllDatafeedsRequest, RequestOptions.DEFAULT);
+        } catch (Exception e1) {
+            logger.warn("failed to stop all datafeeds. Forcing stop", e1);
+            try {
+                stopAllDatafeedsRequest.setForce(true);
+                mlClient.stopDatafeed(stopAllDatafeedsRequest, RequestOptions.DEFAULT);
+            } 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);
+        }
+    }
+
+    private void deleteAllJobs() throws IOException {
+        closeAllJobs();
+
+        GetJobResponse getJobResponse = mlClient.getJob(GetJobRequest.getAllJobsRequest(), RequestOptions.DEFAULT);
+        for (Job job : getJobResponse.jobs()) {
+            mlClient.deleteJob(new DeleteJobRequest(job.getId()), RequestOptions.DEFAULT);
+        }
+    }
+
+    private void closeAllJobs() {
+        CloseJobRequest closeAllJobsRequest = CloseJobRequest.closeAllJobsRequest();
+        try {
+            mlClient.closeJob(closeAllJobsRequest, RequestOptions.DEFAULT);
+        } catch (Exception e1) {
+            logger.warn("failed to close all jobs. Forcing closed", e1);
+            closeAllJobsRequest.setForce(true);
+            try {
+                mlClient.closeJob(closeAllJobsRequest, RequestOptions.DEFAULT);
+            } 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);
+        }
+    }
+}

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

@@ -30,7 +30,7 @@ import org.elasticsearch.action.support.master.AcknowledgedResponse;
 import org.elasticsearch.client.ESRestHighLevelClientTestCase;
 import org.elasticsearch.client.MachineLearningGetResultsIT;
 import org.elasticsearch.client.MachineLearningIT;
-import org.elasticsearch.client.MlRestTestStateCleaner;
+import org.elasticsearch.client.MlTestStateCleaner;
 import org.elasticsearch.client.RequestOptions;
 import org.elasticsearch.client.RestHighLevelClient;
 import org.elasticsearch.client.ml.CloseJobRequest;
@@ -126,7 +126,7 @@ public class MlClientDocumentationIT extends ESRestHighLevelClientTestCase {
 
     @After
     public void cleanUp() throws IOException {
-        new MlRestTestStateCleaner(logger, client()).clearMlMetadata();
+        new MlTestStateCleaner(logger, highLevelClient().machineLearning()).clearMlMetadata();
     }
 
     public void testCreateJob() throws Exception {