Fix integration tests

Original commit: elastic/x-pack-elasticsearch@a5955cfde3

elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/PrelertJobIT.java

@@ -22,7 +22,6 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.stream.Collectors;
 
-import static org.elasticsearch.xpack.prelert.integration.ScheduledJobIT.clearPrelertMetadata;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.isEmptyString;
@@ -286,6 +285,6 @@ public class PrelertJobIT extends ESRestTestCase {
 
     @After
     public void clearPrelertState() throws IOException {
-        clearPrelertMetadata(adminClient());
+        new PrelertRestTestStateCleaner(client(), this).clearPrelertMetadata();
     }
 }

elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/PrelertRestTestStateCleaner.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.prelert.integration;
+
+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.Collections;
+import java.util.List;
+import java.util.Map;
+
+public class PrelertRestTestStateCleaner {
+
+    private final RestClient client;
+    private final ESRestTestCase testCase;
+
+    public PrelertRestTestStateCleaner(RestClient client, ESRestTestCase testCase) {
+        this.client = client;
+        this.testCase = testCase;
+    }
+
+    public void clearPrelertMetadata() throws IOException {
+        deleteAllSchedulers();
+        deleteAllJobs();
+    }
+
+    @SuppressWarnings("unchecked")
+    private void deleteAllSchedulers() throws IOException {
+        Map<String, Object> clusterStateAsMap = testCase.entityAsMap(client.performRequest("GET", "/_cluster/state",
+                Collections.singletonMap("filter_path", "metadata.prelert.schedulers")));
+        List<Map<String, Object>> schedulers =
+                (List<Map<String, Object>>) XContentMapValues.extractValue("metadata.prelert.schedulers", clusterStateAsMap);
+        if (schedulers == null) {
+            return;
+        }
+
+        for (Map<String, Object> scheduler : schedulers) {
+            Map<String, Object> schedulerMap = (Map<String, Object>) scheduler.get("config");
+            String schedulerId = (String) schedulerMap.get("scheduler_id");
+            try {
+                client.performRequest("POST", "/_xpack/prelert/schedulers/" + schedulerId + "/_stop");
+            } catch (Exception e) {
+                // ignore
+            }
+            client.performRequest("DELETE", "/_xpack/prelert/schedulers/" + schedulerId);
+        }
+    }
+
+    private void deleteAllJobs() throws IOException {
+        Map<String, Object> clusterStateAsMap = testCase.entityAsMap(client.performRequest("GET", "/_cluster/state",
+                Collections.singletonMap("filter_path", "metadata.prelert.jobs")));
+        @SuppressWarnings("unchecked")
+        List<Map<String, Object>> jobConfigs =
+                (List<Map<String, Object>>) XContentMapValues.extractValue("metadata.prelert.jobs", clusterStateAsMap);
+        if (jobConfigs == null) {
+            return;
+        }
+
+        for (Map<String, Object> jobConfig : jobConfigs) {
+            String jobId = (String) jobConfig.get("job_id");
+            try {
+                client.performRequest("POST", "/_xpack/prelert/anomaly_detectors/" + jobId + "/_close");
+            } catch (Exception e) {
+                // ignore
+            }
+            client.performRequest("DELETE", "/_xpack/prelert/anomaly_detectors/" + jobId);
+        }
+    }
+}

elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/PrelertYamlTestSuiteIT.java

@@ -5,17 +5,14 @@
  */
 package org.elasticsearch.xpack.prelert.integration;
 
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
 import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
 import org.elasticsearch.test.rest.yaml.parser.ClientYamlTestParseException;
 import org.junit.After;
 
-import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
-
 import java.io.IOException;
 
-import static org.elasticsearch.xpack.prelert.integration.ScheduledJobIT.clearPrelertMetadata;
-
 /** Rest integration test. Runs against a cluster started by {@code gradle integTest} */
 public class PrelertYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
 
@@ -30,7 +27,6 @@ public class PrelertYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
 
     @After
     public void clearPrelertState() throws IOException {
-        clearPrelertMetadata(adminClient());
+        new PrelertRestTestStateCleaner(client(), this).clearPrelertMetadata();
     }
-
 }

elasticsearch/src/test/java/org/elasticsearch/xpack/prelert/integration/ScheduledJobIT.java

@@ -9,7 +9,6 @@ import org.apache.http.entity.StringEntity;
 import org.elasticsearch.client.Response;
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.client.RestClient;
-import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xpack.prelert.PrelertPlugin;
 import org.junit.After;
@@ -19,8 +18,6 @@ import java.io.IOException;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
-import java.util.List;
-import java.util.Map;
 import java.util.stream.Collectors;
 
 import static org.hamcrest.Matchers.containsString;
@@ -134,62 +131,13 @@ public class ScheduledJobIT extends ESRestTestCase {
         }
     }
 
-    @After
-    public void clearPrelertState() throws IOException {
-        clearPrelertMetadata(adminClient());
-    }
-
-    public static void clearPrelertMetadata(RestClient client) throws IOException {
-        deleteAllSchedulers(client);
-        deleteAllJobs(client);
-    }
-
-    @SuppressWarnings("unchecked")
-    private static void deleteAllSchedulers(RestClient client) throws IOException {
-        Map<String, Object> clusterStateAsMap = entityAsMap(client.performRequest("GET", "/_cluster/state",
-                Collections.singletonMap("filter_path", "metadata.prelert.schedulers")));
-        List<Map<String, Object>> schedulers =
-                (List<Map<String, Object>>) XContentMapValues.extractValue("metadata.prelert.schedulers", clusterStateAsMap);
-        if (schedulers == null) {
-            return;
-        }
-
-        for (Map<String, Object> scheduler : schedulers) {
-            Map<String, Object> schedulerMap = (Map<String, Object>) scheduler.get("config");
-            String schedulerId = (String) schedulerMap.get("scheduler_id");
-            try {
-                client.performRequest("POST", "/_xpack/prelert/schedulers/" + schedulerId + "/_stop");
-            } catch (Exception e) {
-                // ignore
-            }
-            client.performRequest("DELETE", "/_xpack/prelert/schedulers/" + schedulerId);
-        }
-    }
-
-    private static void deleteAllJobs(RestClient client) throws IOException {
-        Map<String, Object> clusterStateAsMap = entityAsMap(client.performRequest("GET", "/_cluster/state",
-                Collections.singletonMap("filter_path", "metadata.prelert.jobs")));
-        @SuppressWarnings("unchecked")
-        List<Map<String, Object>> jobConfigs =
-                (List<Map<String, Object>>) XContentMapValues.extractValue("metadata.prelert.jobs", clusterStateAsMap);
-        if (jobConfigs == null) {
-            return;
-        }
-
-        for (Map<String, Object> jobConfig : jobConfigs) {
-            String jobId = (String) jobConfig.get("job_id");
-            try {
-                Response response = client.performRequest("POST", "/_xpack/prelert/anomaly_detectors/" + jobId + "/_close");
-                assertThat(response.getStatusLine().getStatusCode(), equalTo(200));
-            } catch (Exception e) {
-                // ignore
-            }
-            client.performRequest("DELETE", "/_xpack/prelert/anomaly_detectors/" + jobId);
-        }
-    }
-
     public static void openJob(RestClient client, String jobId) throws IOException {
         Response response = client.performRequest("post", PrelertPlugin.BASE_PATH + "anomaly_detectors/" + jobId + "/_open");
         assertThat(response.getStatusLine().getStatusCode(), equalTo(200));
     }
+
+    @After
+    public void clearPrelertState() throws IOException {
+        new PrelertRestTestStateCleaner(client(), this).clearPrelertMetadata();
+    }
 }