From f1d6b33eccdabb3d7ede33a89eeaae804ff090e8 Mon Sep 17 00:00:00 2001
From: Aurelien FOUCRET <aurelien.foucret@elastic.co>
Date: Wed, 20 Mar 2024 08:38:27 +0100
Subject: [PATCH] Ensure ILM policy is installed before starting the tests.

---
 .../application/FullClusterRestartIT.java     | 41 ++++++++++++++++---
 1 file changed, 35 insertions(+), 6 deletions(-)

diff --git a/x-pack/plugin/ent-search/qa/full-cluster-restart/src/javaRestTest/java/org/elasticsearch/xpack/application/FullClusterRestartIT.java b/x-pack/plugin/ent-search/qa/full-cluster-restart/src/javaRestTest/java/org/elasticsearch/xpack/application/FullClusterRestartIT.java
index 95b3b576eb46c..cdbe4d082287b 100644
--- a/x-pack/plugin/ent-search/qa/full-cluster-restart/src/javaRestTest/java/org/elasticsearch/xpack/application/FullClusterRestartIT.java
+++ b/x-pack/plugin/ent-search/qa/full-cluster-restart/src/javaRestTest/java/org/elasticsearch/xpack/application/FullClusterRestartIT.java
@@ -66,10 +66,10 @@ public void testBehavioralAnalyticsDataRetention() throws Exception {
         String legacyAnalyticsCollectionName = "oldstuff";
         String newAnalyticsCollectionName = "newstuff";
 
-        if (isRunningAgainstOldCluster()) {
-            // Ensure index template is installed before executing the tests.
-            assertBusy(() -> assertDataStreamTemplateExists(EVENT_DATA_STREAM_LEGACY_TEMPLATE_NAME));
+        // Wait for the cluster to finish initialization
+        waitForClusterReady();
 
+        if (isRunningAgainstOldCluster()) {
             // Create an analytics collection
             Request legacyPutRequest = new Request("PUT", "_application/analytics/" + legacyAnalyticsCollectionName);
             assertOK(client().performRequest(legacyPutRequest));
@@ -77,9 +77,6 @@ public void testBehavioralAnalyticsDataRetention() throws Exception {
             // Validate that ILM lifecycle is in place
             assertBusy(() -> assertUsingLegacyDataRetentionPolicy(legacyAnalyticsCollectionName));
         } else {
-            // Ensure index template is updated to version 3 before executing the tests.
-            assertBusy(() -> assertDataStreamTemplateExists(EVENT_DATA_STREAM_LEGACY_TEMPLATE_NAME, DSL_REGISTRY_VERSION));
-
             // Create a new analytics collection
             Request putRequest = new Request("PUT", "_application/analytics/" + newAnalyticsCollectionName);
             assertOK(client().performRequest(putRequest));
@@ -129,6 +126,21 @@ private void assertUsingDslDataRetention(String analyticsCollectionName) throws
         assertTrue(evaluatedNewDataStream);
     }
 
+    private void waitForClusterReady() throws Exception {
+        // Ensure index template is installed with the right version before executing the tests.
+        if (isRunningAgainstOldCluster()) {
+            // No minimum version of the registry required when running on old clusters.
+            assertBusy(() -> assertDataStreamTemplateExists(EVENT_DATA_STREAM_LEGACY_TEMPLATE_NAME));
+
+            // When running on old cluster, wait for the ILM policy to be installed.
+            assertBusy(() -> assertILMPolicyExists(EVENT_DATA_STREAM_LEGACY_ILM_POLICY_NAME));
+        } else {
+            // DSL has been introduced with the version 3 of the registry.
+            // Wait for this version to be deployed.
+            assertBusy(() -> assertDataStreamTemplateExists(EVENT_DATA_STREAM_LEGACY_TEMPLATE_NAME, DSL_REGISTRY_VERSION));
+        }
+    }
+
     private void assertDataStreamTemplateExists(String templateName) throws IOException {
         assertDataStreamTemplateExists(templateName, null);
     }
@@ -138,6 +150,7 @@ private void assertDataStreamTemplateExists(String templateName, Integer minVers
             Request getIndexTemplateRequest = new Request("GET", "_index_template/" + templateName);
             Response response = client().performRequest(getIndexTemplateRequest);
             assertOK(response);
+
             if (minVersion != null) {
                 String pathToVersion = "index_templates.0.index_template.version";
                 ObjectPath indexTemplatesResponse = ObjectPath.createFromResponse(response);
@@ -151,4 +164,20 @@ private void assertDataStreamTemplateExists(String templateName, Integer minVers
             throw e;
         }
     }
+
+    private void assertILMPolicyExists(String policyName) throws IOException {
+        try {
+            Request getILMPolicyRequest = new Request("GET", "_ilm/policy/" + policyName);
+            Response response = client().performRequest(getILMPolicyRequest);
+            assertOK(response);
+
+            assertNotNull(ObjectPath.createFromResponse(response).evaluate(policyName));
+        }  catch (ResponseException e) {
+            int status = e.getResponse().getStatusLine().getStatusCode();
+            if (status == 404) {
+                throw new AssertionError("Waiting for the policy to be created");
+            }
+            throw e;
+        }
+    }
 }