[TESTS] Refactor yml test suite classes (elastic/x-pack-elasticsearch#3145)

* Use XPackRestIT as base class for XDocsClientYamlTestSuiteIT * Remove the XPackRestTestCase class * Address review comments * Fix checkstyle checks Original commit: elastic/x-pack-elasticsearch@c2a5e60c12
2025-03-24 17:09:48 +00:00 · 2017-11-29 12:43:53 +00:00 · 2017-11-29 12:43:53 +00:00 · 171c48fd2f
commit 171c48fd2f
parent ef96831515
10 changed files with 256 additions and 275 deletions
--- a/docs/src/test/java/org/elasticsearch/smoketest/XDocsClientYamlTestSuiteIT.java
+++ b/docs/src/test/java/org/elasticsearch/smoketest/XDocsClientYamlTestSuiteIT.java
@ -7,7 +7,6 @@ package org.elasticsearch.smoketest;

 import org.apache.http.HttpHost;
 import com.carrotsearch.randomizedtesting.annotations.Name;
-import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 import org.elasticsearch.Version;
 import org.elasticsearch.client.RestClient;
 import org.elasticsearch.common.settings.SecureString;
@ -17,9 +16,8 @@ import org.elasticsearch.test.rest.yaml.ClientYamlDocsTestClient;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestClient;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestResponse;
-import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
 import org.elasticsearch.test.rest.yaml.restspec.ClientYamlSuiteRestSpec;
-import org.elasticsearch.xpack.ml.integration.MlRestTestStateCleaner;
+import org.elasticsearch.xpack.test.rest.XPackRestIT;
 import org.junit.After;

 import java.io.IOException;
@ -32,18 +30,13 @@ import static java.util.Collections.singletonMap;
 import static org.elasticsearch.xpack.security.authc.support.UsernamePasswordToken.basicAuthHeaderValue;
 import static org.hamcrest.Matchers.is;

-public class XDocsClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
+public class XDocsClientYamlTestSuiteIT extends XPackRestIT {
    private static final String USER_TOKEN = basicAuthHeaderValue("test_admin", new SecureString("x-pack-test-password".toCharArray()));

    public XDocsClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {
        super(testCandidate);
    }

-    @ParametersFactory
-    public static Iterable<Object[]> parameters() throws Exception {
-        return ESClientYamlSuiteTestCase.createParameters();
-    }
-
    @Override
    protected void afterIfFailed(List<Throwable> errors) {
        super.afterIfFailed(errors);
@ -93,11 +86,23 @@ public class XDocsClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
        }
    }

-    private boolean isWatcherTest() {
+    @Override
+    protected boolean isWatcherTest() {
        String testName = getTestName();
        return testName != null && testName.contains("watcher");
    }

+    @Override
+    protected boolean isMonitoringTest() {
+        return false;
+    }
+
+    @Override
+    protected boolean isMachineLearningTest() {
+        String testName = getTestName();
+        return testName != null && testName.contains("ml/");
+    }
+
    /**
     * Deletes users after every test just in case any test adds any.
     */
@ -117,11 +122,6 @@ public class XDocsClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {
        }
    }

-    @After
-    public void cleanMlState() throws Exception {
-        new MlRestTestStateCleaner(logger, adminClient(), this).clearMlMetadata();
-    }
-
    @Override
    protected boolean randomizeContentType() {
        return false;
--- a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/DatafeedJobsRestIT.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/DatafeedJobsRestIT.java
@ -16,6 +16,7 @@ import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.test.SecuritySettingsSource;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xpack.ml.MachineLearning;
+import org.elasticsearch.xpack.test.rest.XPackRestTestHelper;
 import org.junit.After;
 import org.junit.Before;

@ -720,6 +721,7 @@ public class DatafeedJobsRestIT extends ESRestTestCase {
    @After
    public void clearMlState() throws Exception {
        new MlRestTestStateCleaner(logger, adminClient(), this).clearMlMetadata();
+        XPackRestTestHelper.waitForPendingTasks(adminClient());
    }

    private static class DatafeedBuilder {
--- a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlJobIT.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlJobIT.java
@ -17,6 +17,7 @@ import org.elasticsearch.test.SecuritySettingsSource;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.xpack.ml.MachineLearning;
 import org.elasticsearch.xpack.ml.job.persistence.AnomalyDetectorsIndex;
+import org.elasticsearch.xpack.test.rest.XPackRestTestHelper;
 import org.junit.After;

 import java.io.BufferedReader;
@ -648,5 +649,6 @@ public class MlJobIT extends ESRestTestCase {
    @After
    public void clearMlState() throws Exception {
        new MlRestTestStateCleaner(logger, adminClient(), this).clearMlMetadata();
+        XPackRestTestHelper.waitForPendingTasks(adminClient());
    }
 }
--- a/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlRestTestStateCleaner.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/ml/integration/MlRestTestStateCleaner.java
@ -37,43 +37,12 @@ public class MlRestTestStateCleaner {
        this.testCase = testCase;
    }

-    public void clearMlMetadata() throws Exception {
+    public void clearMlMetadata() throws IOException {
        deleteAllDatafeeds();
        deleteAllJobs();
-        waitForPendingTasks();
        // indices will be deleted by the ESIntegTestCase class
    }

-    private void waitForPendingTasks() throws Exception {
-        ESTestCase.assertBusy(() -> {
-            try {
-                Response response = adminClient.performRequest("GET", "/_cat/tasks",
-                        Collections.singletonMap("detailed", "true"));
-                // Check to see if there are tasks still active. We exclude the
-                // list tasks
-                // actions tasks form this otherwise we will always fail
-                if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
-                    try (BufferedReader responseReader = new BufferedReader(
-                            new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8))) {
-                        int activeTasks = 0;
-                        String line;
-                        StringBuilder tasksListString = new StringBuilder();
-                        while ((line = responseReader.readLine()) != null) {
-                            if (line.startsWith(ListTasksAction.NAME) == false) {
-                                activeTasks++;
-                                tasksListString.append(line);
-                                tasksListString.append('\n');
-                            }
-                        }
-                        assertEquals(activeTasks + " active tasks found:\n" + tasksListString, 0, activeTasks);
-                    }
-                }
-            } catch (IOException e) {
-                throw new AssertionError("Error getting active tasks list", e);
-            }
-        });
-    }
-
    @SuppressWarnings("unchecked")
    private void deleteAllDatafeeds() throws IOException {
        Map<String, Object> clusterStateAsMap = testCase.entityAsMap(adminClient.performRequest("GET", "/_cluster/state",
--- a/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestIT.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestIT.java
@ -5,13 +5,18 @@
 */
 package org.elasticsearch.xpack.test.rest;

+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
 import org.apache.http.HttpStatus;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.common.CheckedFunction;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.util.concurrent.ThreadContext;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.plugins.MetaDataUpgrader;
+import org.elasticsearch.test.SecuritySettingsSource;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestResponse;
+import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
 import org.elasticsearch.xpack.ml.MlMetaIndex;
 import org.elasticsearch.xpack.ml.integration.MlRestTestStateCleaner;
 import org.elasticsearch.xpack.ml.job.persistence.AnomalyDetectorsIndex;
@ -24,6 +29,7 @@ import org.junit.Before;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@ -35,33 +41,60 @@ import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonList;
 import static java.util.Collections.singletonMap;
 import static org.elasticsearch.common.xcontent.support.XContentMapValues.extractValue;
+import static org.elasticsearch.xpack.security.authc.support.UsernamePasswordToken.basicAuthHeaderValue;
 import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.is;

 /** Runs rest tests against external cluster */
-public class XPackRestIT extends XPackRestTestCase {
+public class XPackRestIT extends ESClientYamlSuiteTestCase {
+
+    private static final String BASIC_AUTH_VALUE =
+            basicAuthHeaderValue("x_pack_rest_user", SecuritySettingsSource.TEST_PASSWORD_SECURE_STRING);

    public XPackRestIT(ClientYamlTestCandidate testCandidate) {
        super(testCandidate);
    }

+    @ParametersFactory
+    public static Iterable<Object[]> parameters() throws Exception {
+        return ESClientYamlSuiteTestCase.createParameters();
+    }
+
+    @Override
+    protected Settings restClientSettings() {
+        return Settings.builder()
+                .put(ThreadContext.PREFIX + ".Authorization", BASIC_AUTH_VALUE)
+                .build();
+    }
+
+
+    @Before
+    public void setupForTests() throws Exception {
+        waitForTemplates();
+        waitForWatcher();
+        enableMonitoring();
+    }
+
    /**
     * Waits for the Security template and the Machine Learning templates to be created by the {@link MetaDataUpgrader}
     */
-    @Before
-    public void waitForTemplates() throws Exception {
-        List<String> templates = new ArrayList<>();
-        templates.add(SecurityLifecycleService.SECURITY_TEMPLATE_NAME);
-        templates.addAll(Arrays.asList(Auditor.NOTIFICATIONS_INDEX, MlMetaIndex.INDEX_NAME, AnomalyDetectorsIndex.jobStateIndexName(),
-                AnomalyDetectorsIndex.jobResultsIndexPrefix()));
-        templates.addAll(Arrays.asList(WatcherIndexTemplateRegistry.TEMPLATE_NAMES));
+    private void waitForTemplates() throws Exception {
+        if (installTemplates()) {
+            List<String> templates = new ArrayList<>();
+            templates.add(SecurityLifecycleService.SECURITY_TEMPLATE_NAME);
+            templates.addAll(Arrays.asList(Auditor.NOTIFICATIONS_INDEX, MlMetaIndex.INDEX_NAME, AnomalyDetectorsIndex.jobStateIndexName(),
+                    AnomalyDetectorsIndex.jobResultsIndexPrefix()));
+            templates.addAll(Arrays.asList(WatcherIndexTemplateRegistry.TEMPLATE_NAMES));

-        for (String template : templates) {
-            awaitCallApi("indices.exists_template", singletonMap("name", template), emptyList(),
-                    response -> true,
-                    () -> "Exception when waiting for [" + template + "] template to be created");
+            for (String template : templates) {
+                awaitCallApi("indices.exists_template", singletonMap("name", template), emptyList(),
+                        response -> true,
+                        () -> "Exception when waiting for [" + template + "] template to be created");
+            }
        }
+    }

+    private void waitForWatcher() throws Exception {
        // ensure watcher is started, so that a test can stop watcher and everything still works fine
        if (isWatcherTest()) {
            assertBusy(() -> {
@ -86,11 +119,10 @@ public class XPackRestIT extends XPackRestTestCase {
     * monitoring indices.This is the signal that the local exporter is started and ready
     * for the tests.
     */
-    @Before
-    public void enableMonitoring() throws Exception {
+    private void enableMonitoring() throws Exception {
        if (isMonitoringTest()) {
            final ClientYamlTestResponse xpackUsage =
-                    callApi("xpack.usage", singletonMap("filter_path", "monitoring.enabled_exporters"), emptyList());
+                    callApi("xpack.usage", singletonMap("filter_path", "monitoring.enabled_exporters"), emptyList(), getApiCallHeaders());

            @SuppressWarnings("unchecked")
            final Map<String, Object> exporters = (Map<String, Object>) xpackUsage.evaluate("monitoring.enabled_exporters");
@ -115,21 +147,6 @@ public class XPackRestIT extends XPackRestTestCase {
        }
    }

-    /**
-     * Cleanup after tests.
-     *
-     * Feature-specific cleanup methods should be called from here rather than using
-     * separate @After annotated methods to ensure there is a well-defined cleanup order.
-     */
-    @After
-    public void cleanup() throws Exception {
-        disableMonitoring();
-        // This also generically waits for pending tasks to complete, so must go last (otherwise
-        // it could be waiting for pending tasks while monitoring is still running).
-        // TODO: consider moving the bit that waits for pending tasks into a general X-Pack component
-        clearMlState();
-    }
-
    /**
     * Disable monitoring
     */
@ -150,7 +167,8 @@ public class XPackRestIT extends XPackRestTestCase {
            awaitBusy(() -> {
                try {
                    ClientYamlTestResponse response =
-                            callApi("xpack.usage", singletonMap("filter_path", "monitoring.enabled_exporters"), emptyList());
+                            callApi("xpack.usage", singletonMap("filter_path", "monitoring.enabled_exporters"), emptyList(),
+                                    getApiCallHeaders());

                    @SuppressWarnings("unchecked")
                    final Map<String, ?> exporters = (Map<String, ?>) response.evaluate("monitoring.enabled_exporters");
@ -162,7 +180,7 @@ public class XPackRestIT extends XPackRestTestCase {
                    params.put("node_id", "_local");
                    params.put("metric", "thread_pool");
                    params.put("filter_path", "nodes.*.thread_pool.bulk.active");
-                    response = callApi("nodes.stats", params, emptyList());
+                    response = callApi("nodes.stats", params, emptyList(), getApiCallHeaders());

                    @SuppressWarnings("unchecked")
                    final Map<String, Object> nodes = (Map<String, Object>) response.evaluate("nodes");
@ -180,10 +198,24 @@ public class XPackRestIT extends XPackRestTestCase {
    }

    /**
-     * Delete any left over machine learning datafeeds and jobs.
+     * Cleanup after tests.
     *
-     * Also waits for pending tasks to complete (which is not really an ML-specific
-     * thing and could be moved into a general X-Pack method at some point).
+     * Feature-specific cleanup methods should be called from here rather than using
+     * separate @After annotated methods to ensure there is a well-defined cleanup order.
+     */
+    @After
+    public void cleanup() throws Exception {
+        disableMonitoring();
+        clearMlState();
+        if (isWaitForPendingTasks()) {
+            // This waits for pending tasks to complete, so must go last (otherwise
+            // it could be waiting for pending tasks while monitoring is still running).
+            XPackRestTestHelper.waitForPendingTasks(adminClient());
+        }
+    }
+
+    /**
+     * Delete any left over machine learning datafeeds and jobs.
     */
    private void clearMlState() throws Exception {
        if (isMachineLearningTest()) {
@ -203,7 +235,7 @@ public class XPackRestIT extends XPackRestTestCase {
        AtomicReference<IOException> exceptionHolder = new AtomicReference<>();
        awaitBusy(() -> {
            try {
-                ClientYamlTestResponse response = callApi(apiName, params, bodies);
+                ClientYamlTestResponse response = callApi(apiName, params, bodies, getApiCallHeaders());
                if (response.getStatusCode() == HttpStatus.SC_OK) {
                    exceptionHolder.set(null);
                    return success.apply(response);
@ -223,22 +255,40 @@ public class XPackRestIT extends XPackRestTestCase {

    private ClientYamlTestResponse callApi(String apiName,
                                           Map<String, String> params,
-                                           List<Map<String, Object>> bodies) throws IOException {
-        return getAdminExecutionContext().callApi(apiName, params, bodies, emptyMap());
+                                           List<Map<String, Object>> bodies,
+                                           Map<String, String> headers) throws IOException {
+        return getAdminExecutionContext().callApi(apiName, params, bodies, headers);
    }

-    private boolean isMonitoringTest() {
+    protected Map<String, String> getApiCallHeaders() {
+        return Collections.emptyMap();
+    }
+
+    protected boolean installTemplates() {
+        return true;
+    }
+
+    protected boolean isMonitoringTest() {
        String testName = getTestName();
        return testName != null && (testName.contains("=monitoring/") || testName.contains("=monitoring\\"));
    }

-    private boolean isWatcherTest() {
+    protected boolean isWatcherTest() {
        String testName = getTestName();
        return testName != null && (testName.contains("=watcher/") || testName.contains("=watcher\\"));
    }

-    private boolean isMachineLearningTest() {
+    protected boolean isMachineLearningTest() {
        String testName = getTestName();
        return testName != null && (testName.contains("=ml/") || testName.contains("=ml\\"));
    }
+
+    /**
+     * Should each test wait for pending tasks to finish after execution?
+     * @return Wait for pending tasks
+     */
+    protected boolean isWaitForPendingTasks() {
+        return true;
+    }
+
 }
--- a/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestTestCase.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestTestCase.java
@ -1,98 +0,0 @@
-/*
- * 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.test.rest;
-
-
-import com.carrotsearch.randomizedtesting.annotations.Name;
-import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
-import org.apache.http.util.EntityUtils;
-import org.elasticsearch.Version;
-import org.elasticsearch.client.ResponseException;
-import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.common.util.concurrent.ThreadContext;
-import org.elasticsearch.common.xcontent.XContentHelper;
-import org.elasticsearch.common.xcontent.json.JsonXContent;
-import org.elasticsearch.test.SecuritySettingsSource;
-import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
-import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
-import org.elasticsearch.xpack.ml.MlMetaIndex;
-import org.elasticsearch.xpack.ml.job.persistence.AnomalyDetectorsIndex;
-import org.elasticsearch.xpack.ml.notifications.Auditor;
-
-import java.io.IOException;
-import java.util.Arrays;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.atomic.AtomicReference;
-
-import static java.util.Collections.singletonMap;
-import static org.elasticsearch.xpack.security.authc.support.UsernamePasswordToken.basicAuthHeaderValue;
-
-public abstract class XPackRestTestCase extends ESClientYamlSuiteTestCase {
-
-    private static final String BASIC_AUTH_VALUE =
-            basicAuthHeaderValue("x_pack_rest_user", SecuritySettingsSource.TEST_PASSWORD_SECURE_STRING);
-
-    public XPackRestTestCase(@Name("yaml") ClientYamlTestCandidate testCandidate) {
-        super(testCandidate);
-    }
-
-    @ParametersFactory
-    public static Iterable<Object[]> parameters() throws Exception {
-        return ESClientYamlSuiteTestCase.createParameters();
-    }
-
-    @Override
-    protected Settings restClientSettings() {
-        return Settings.builder()
-                .put(ThreadContext.PREFIX + ".Authorization", BASIC_AUTH_VALUE)
-                .build();
-    }
-
-    /**
-     * Waits for the Machine Learning templates to be created by {@link org.elasticsearch.plugins.MetaDataUpgrader}.
-     */
-    public static void waitForMlTemplates() throws InterruptedException {
-        AtomicReference<Version> masterNodeVersion = new AtomicReference<>();
-        awaitBusy(() -> {
-            String response;
-            try {
-                response = EntityUtils
-                        .toString(client().performRequest("GET", "/_cat/nodes", singletonMap("h", "master,version")).getEntity());
-            } catch (IOException e) {
-                throw new RuntimeException(e);
-            }
-            for (String line : response.split("\n")) {
-                if (line.startsWith("*")) {
-                    masterNodeVersion.set(Version.fromString(line.substring(2).trim()));
-                    return true;
-                }
-            }
-            return false;
-        });
-
-        final List<String> templateNames = Arrays.asList(Auditor.NOTIFICATIONS_INDEX, MlMetaIndex.INDEX_NAME,
-                AnomalyDetectorsIndex.jobStateIndexName(), AnomalyDetectorsIndex.jobResultsIndexPrefix());
-        for (String template : templateNames) {
-            awaitBusy(() -> {
-                Map<?, ?> response;
-                try {
-                    String string = EntityUtils.toString(client().performRequest("GET", "/_template/" + template).getEntity());
-                    response = XContentHelper.convertToMap(JsonXContent.jsonXContent, string, false);
-                } catch (ResponseException e) {
-                    if (e.getResponse().getStatusLine().getStatusCode() == 404) {
-                        return false;
-                    }
-                    throw new RuntimeException(e);
-                } catch (IOException e) {
-                    throw new RuntimeException(e);
-                }
-                Map<?, ?> templateDefinition = (Map<?, ?>) response.get(template);
-                return Version.fromId((Integer) templateDefinition.get("version")).equals(masterNodeVersion.get());
-            });
-        }
-    }
-}
--- a/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestTestHelper.java
+++ b/plugin/src/test/java/org/elasticsearch/xpack/test/rest/XPackRestTestHelper.java
@ -0,0 +1,118 @@
+/*
+ * 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.test.rest;
+
+
+import org.apache.http.HttpStatus;
+import org.apache.http.util.EntityUtils;
+import org.elasticsearch.Version;
+import org.elasticsearch.action.admin.cluster.node.tasks.list.ListTasksAction;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.ResponseException;
+import org.elasticsearch.client.RestClient;
+import org.elasticsearch.common.xcontent.XContentHelper;
+import org.elasticsearch.common.xcontent.json.JsonXContent;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.ml.MlMetaIndex;
+import org.elasticsearch.xpack.ml.job.persistence.AnomalyDetectorsIndex;
+import org.elasticsearch.xpack.ml.notifications.Auditor;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static java.util.Collections.singletonMap;
+import static org.junit.Assert.assertEquals;
+
+public final class XPackRestTestHelper {
+
+    private XPackRestTestHelper() {
+    }
+
+    /**
+     * Waits for the Machine Learning templates to be created
+     * and check the version is up to date
+     */
+    public static void waitForMlTemplates(RestClient client) throws InterruptedException {
+        AtomicReference<Version> masterNodeVersion = new AtomicReference<>();
+        ESTestCase.awaitBusy(() -> {
+            String response;
+            try {
+                response = EntityUtils
+                        .toString(client.performRequest("GET", "/_cat/nodes", singletonMap("h", "master,version")).getEntity());
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+            for (String line : response.split("\n")) {
+                if (line.startsWith("*")) {
+                    masterNodeVersion.set(Version.fromString(line.substring(2).trim()));
+                    return true;
+                }
+            }
+            return false;
+        });
+
+        final List<String> templateNames = Arrays.asList(Auditor.NOTIFICATIONS_INDEX, MlMetaIndex.INDEX_NAME,
+                AnomalyDetectorsIndex.jobStateIndexName(), AnomalyDetectorsIndex.jobResultsIndexPrefix());
+        for (String template : templateNames) {
+            ESTestCase.awaitBusy(() -> {
+                Map<?, ?> response;
+                try {
+                    String string = EntityUtils.toString(client.performRequest("GET", "/_template/" + template).getEntity());
+                    response = XContentHelper.convertToMap(JsonXContent.jsonXContent, string, false);
+                } catch (ResponseException e) {
+                    if (e.getResponse().getStatusLine().getStatusCode() == 404) {
+                        return false;
+                    }
+                    throw new RuntimeException(e);
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+                Map<?, ?> templateDefinition = (Map<?, ?>) response.get(template);
+                return Version.fromId((Integer) templateDefinition.get("version")).equals(masterNodeVersion.get());
+            });
+        }
+    }
+
+    /**
+     * Waits for pending tasks to complete
+     */
+    public static void waitForPendingTasks(RestClient adminClient) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            try {
+                Response response = adminClient.performRequest("GET", "/_cat/tasks",
+                        Collections.singletonMap("detailed", "true"));
+                // Check to see if there are tasks still active. We exclude the
+                // list tasks
+                // actions tasks form this otherwise we will always fail
+                if (response.getStatusLine().getStatusCode() == HttpStatus.SC_OK) {
+                    try (BufferedReader responseReader = new BufferedReader(
+                            new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8))) {
+                        int activeTasks = 0;
+                        String line;
+                        StringBuilder tasksListString = new StringBuilder();
+                        while ((line = responseReader.readLine()) != null) {
+                            if (line.startsWith(ListTasksAction.NAME) == false) {
+                                activeTasks++;
+                                tasksListString.append(line);
+                                tasksListString.append('\n');
+                            }
+                        }
+                        assertEquals(activeTasks + " active tasks found:\n" + tasksListString, 0, activeTasks);
+                    }
+                }
+            } catch (IOException e) {
+                throw new AssertionError("Error getting active tasks list", e);
+            }
+        });
+    }
+}
--- a/qa/full-cluster-restart/src/test/java/org/elasticsearch/xpack/restart/FullClusterRestartIT.java
+++ b/qa/full-cluster-restart/src/test/java/org/elasticsearch/xpack/restart/FullClusterRestartIT.java
@ -23,7 +23,7 @@ import org.elasticsearch.xpack.watcher.common.text.TextTemplate;
 import org.elasticsearch.xpack.monitoring.exporter.MonitoringTemplateUtils;
 import org.elasticsearch.xpack.security.SecurityClusterClientYamlTestCase;
 import org.elasticsearch.xpack.security.support.IndexLifecycleManager;
-import org.elasticsearch.xpack.test.rest.XPackRestTestCase;
+import org.elasticsearch.xpack.test.rest.XPackRestTestHelper;
 import org.elasticsearch.xpack.watcher.actions.logging.LoggingAction;
 import org.elasticsearch.xpack.watcher.client.WatchSourceBuilder;
 import org.elasticsearch.xpack.watcher.condition.AlwaysCondition;
@ -68,7 +68,7 @@ public class FullClusterRestartIT extends ESRestTestCase {

    @Before
    public void waitForMlTemplates() throws Exception {
-        XPackRestTestCase.waitForMlTemplates();
+        XPackRestTestHelper.waitForMlTemplates(client());
    }

    @Override
--- a/qa/rolling-upgrade/src/test/java/org/elasticsearch/upgrades/UpgradeClusterClientYamlTestSuiteIT.java
+++ b/qa/rolling-upgrade/src/test/java/org/elasticsearch/upgrades/UpgradeClusterClientYamlTestSuiteIT.java
@ -15,7 +15,7 @@ import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestResponse;
 import org.elasticsearch.xpack.security.SecurityClusterClientYamlTestCase;
-import org.elasticsearch.xpack.test.rest.XPackRestTestCase;
+import org.elasticsearch.xpack.test.rest.XPackRestTestHelper;
 import org.junit.Before;

 import java.nio.charset.StandardCharsets;
@ -36,7 +36,7 @@ public class UpgradeClusterClientYamlTestSuiteIT extends SecurityClusterClientYa
     */
    @Before
    public void waitForTemplates() throws Exception {
-        XPackRestTestCase.waitForMlTemplates();
+        XPackRestTestHelper.waitForMlTemplates(client());
    }

    /**
--- a/qa/smoke-test-ml-with-security/src/test/java/org/elasticsearch/smoketest/MlWithSecurityIT.java
+++ b/qa/smoke-test-ml-with-security/src/test/java/org/elasticsearch/smoketest/MlWithSecurityIT.java
@ -6,111 +6,29 @@
 package org.elasticsearch.smoketest;

 import com.carrotsearch.randomizedtesting.annotations.Name;
-import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;

-import org.apache.http.HttpStatus;
-import org.elasticsearch.common.CheckedFunction;
 import org.elasticsearch.common.settings.SecureString;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.ThreadContext;
-import org.elasticsearch.plugins.MetaDataUpgrader;
 import org.elasticsearch.test.SecuritySettingsSource;
 import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
-import org.elasticsearch.test.rest.yaml.ClientYamlTestResponse;
-import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
-import org.elasticsearch.xpack.ml.MlMetaIndex;
-import org.elasticsearch.xpack.ml.integration.MlRestTestStateCleaner;
-import org.elasticsearch.xpack.ml.job.persistence.AnomalyDetectorsIndex;
-import org.elasticsearch.xpack.ml.notifications.Auditor;
-import org.elasticsearch.xpack.security.SecurityLifecycleService;
-import org.junit.After;
-import org.junit.Before;
+import org.elasticsearch.xpack.test.rest.XPackRestIT;
+

-import java.io.IOException;
-import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collections;
-import java.util.List;
 import java.util.Map;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Supplier;

 import static org.elasticsearch.xpack.security.authc.support.UsernamePasswordToken.basicAuthHeaderValue;


-public class MlWithSecurityIT extends ESClientYamlSuiteTestCase {
+public class MlWithSecurityIT extends XPackRestIT {

    private static final String TEST_ADMIN_USERNAME = "x_pack_rest_user";

-    @After
-    public void clearMlState() throws Exception {
-        new MlRestTestStateCleaner(logger, adminClient(), this).clearMlMetadata();
-    }
-
    public MlWithSecurityIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {
        super(testCandidate);
    }

-    /**
-     * Waits for the Security template and the Machine Learning templates to be created by the {@link MetaDataUpgrader}
-     */
-    @Before
-    public void waitForTemplates() throws Exception {
-        List<String> templates = new ArrayList<>();
-        templates.add(SecurityLifecycleService.SECURITY_TEMPLATE_NAME);
-        templates.addAll(Arrays.asList(Auditor.NOTIFICATIONS_INDEX, MlMetaIndex.INDEX_NAME,
-                AnomalyDetectorsIndex.jobStateIndexName(), AnomalyDetectorsIndex.jobResultsIndexPrefix()));
-
-        for (String template : templates) {
-            awaitCallApi("indices.exists_template", Collections.singletonMap("name", template), Collections.emptyList(),
-                    response -> true,
-                    () -> "Exception when waiting for [" + template + "] template to be created");
-        }
-    }
-
-    /**
-     * Executes an API call using the admin context, waiting for it to succeed.
-     */
-    private void awaitCallApi(String apiName,
-                              Map<String, String> params,
-                              List<Map<String, Object>> bodies,
-                              CheckedFunction<ClientYamlTestResponse, Boolean, IOException> success,
-                              Supplier<String> error) throws Exception {
-
-        AtomicReference<IOException> exceptionHolder = new AtomicReference<>();
-        awaitBusy(() -> {
-            try {
-                ClientYamlTestResponse response = callApi(apiName, params, bodies);
-                if (response.getStatusCode() == HttpStatus.SC_OK) {
-                    exceptionHolder.set(null);
-                    return success.apply(response);
-                }
-                return false;
-            } catch (IOException e) {
-                exceptionHolder.set(e);
-            }
-            return false;
-        });
-
-        IOException exception = exceptionHolder.get();
-        if (exception != null) {
-            throw new IllegalStateException(error.get(), exception);
-        }
-    }
-
-    private ClientYamlTestResponse callApi(String apiName,
-                                           Map<String, String> params,
-                                           List<Map<String, Object>> bodies) throws IOException {
-        Map<String, String> headers = Collections.singletonMap("Authorization",
-                basicAuthHeaderValue(TEST_ADMIN_USERNAME, SecuritySettingsSource.TEST_PASSWORD_SECURE_STRING));
-        return getAdminExecutionContext().callApi(apiName, params, bodies, headers);
-    }
-
-    @ParametersFactory
-    public static Iterable<Object[]> parameters() throws Exception {
-        return ESClientYamlSuiteTestCase.createParameters();
-    }
-
    protected String[] getCredentials() {
        return new String[]{"ml_admin", "x-pack-test-password"};
    }
@ -131,4 +49,24 @@ public class MlWithSecurityIT extends ESClientYamlSuiteTestCase {
            .put(ThreadContext.PREFIX + ".Authorization", token)
            .build();
    }
+
+    protected Map<String, String> getApiCallHeaders() {
+        return Collections.singletonMap("Authorization", basicAuthHeaderValue(TEST_ADMIN_USERNAME,
+                SecuritySettingsSource.TEST_PASSWORD_SECURE_STRING));
+    }
+
+    @Override
+    protected boolean isMonitoringTest() {
+        return false;
+    }
+
+    @Override
+    protected boolean isWatcherTest() {
+        return false;
+    }
+
+    @Override
+    protected boolean isMachineLearningTest() {
+        return true;
+    }
 }