[ML][HLRC] Replace REST-based ML test cleanup with the ML client

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

@@ -147,7 +147,7 @@ private void addCategoriesIndexRequests(BulkRequest bulkRequest) {
 
     @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/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.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 {