/*
 * 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.upgrades;

import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
import com.carrotsearch.randomizedtesting.annotations.TimeoutSuite;
import org.apache.lucene.util.TimeUnits;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.util.concurrent.ThreadContext;
import org.elasticsearch.test.rest.ESRestTestCase;
import org.elasticsearch.test.rest.yaml.ClientYamlTestCandidate;
import org.elasticsearch.test.rest.yaml.ESClientYamlSuiteTestCase;
import org.elasticsearch.xpack.test.rest.XPackRestTestConstants;
import org.elasticsearch.xpack.test.rest.XPackRestTestHelper;
import org.junit.Before;

import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;

@TimeoutSuite(millis = 5 * TimeUnits.MINUTE) // to account for slow as hell VMs
public class UpgradeClusterClientYamlTestSuiteIT extends ESClientYamlSuiteTestCase {

    /**
     * Waits for the Machine Learning templates to be created by {@link org.elasticsearch.plugins.MetadataUpgrader}
     */
    @Before
    public void waitForTemplates() throws Exception {
        XPackRestTestHelper.waitForTemplates(client(), XPackRestTestConstants.ML_POST_V660_TEMPLATES);
    }

    @Before
    public void waitForWatcher() throws Exception {
        // Wait for watcher to be in started state in order to avoid errors due
        // to manually executing watches prior for watcher to be ready:
        assertBusy(() -> {
            Response response = client().performRequest(new Request("GET", "_watcher/stats"));
            Map<String, Object> responseBody = entityAsMap(response);
            List<?> stats = (List<?>) responseBody.get("stats");
            assertThat(stats.size(), greaterThanOrEqualTo(3));
            for (Object stat : stats) {
                Map<?, ?> statAsMap = (Map<?, ?>) stat;
                assertThat(statAsMap.get("watcher_state"), equalTo("started"));
            }
        }, 1, TimeUnit.MINUTES);
    }

    @Override
    protected boolean preserveIndicesUponCompletion() {
        return true;
    }

    @Override
    protected boolean preserveTemplatesUponCompletion() {
        return true;
    }

    @Override
    protected boolean preserveRollupJobsUponCompletion() {
        return true;
    }

    @Override
    protected boolean preserveILMPoliciesUponCompletion() {
        return true;
    }

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

    @ParametersFactory
    public static Iterable<Object[]> parameters() throws Exception {
        return createParameters();
    }

    @Override
    protected Settings restClientSettings() {
        String token = "Basic " + Base64.getEncoder().encodeToString(("test_user:x-pack-test-password").getBytes(StandardCharsets.UTF_8));
        return Settings.builder()
                .put(ThreadContext.PREFIX + ".Authorization", token)
                // we increase the timeout here to 90 seconds to handle long waits for a green
                // cluster health. the waits for green need to be longer than a minute to
                // account for delayed shards
                .put(ESRestTestCase.CLIENT_SOCKET_TIMEOUT, "90s")
                .build();
    }
}

@@ -1598,32 +1598,50 @@
 ows Exception %7B%0A
+        try %7B%0A    
         XPackRes
@@ -1723,24 +1723,178 @@
 TEMPLATES);%0A
+        %7D catch (AssertionError e) %7B%0A            throw new AssertionError(%22Failure in test setup: Failed to initialize ML index templates%22, e);%0A        %7D%0A
     %7D%0A%0A    @
@@ -2104,16 +2104,34 @@
  ready:%0A
+        try %7B%0A    
         
@@ -2161,16 +2161,20 @@
         
+   
+ 
 Response
@@ -2256,24 +2256,28 @@
             
+    
 Map%3CString, 
@@ -2334,16 +2334,20 @@
         
+  
+  
 List%3C?%3E 
@@ -2391,16 +2391,20 @@
 tats%22);%0A
+    
         
@@ -2470,16 +2470,20 @@
         
+    
 for (Obj
@@ -2498,24 +2498,28 @@
  : stats) %7B%0A
+    
             
@@ -2562,16 +2562,20 @@
 ) stat;%0A
+    
         
@@ -2654,26 +2654,34 @@
             
-%7D%0A
+    %7D%0A    
         %7D, 1
@@ -2697,24 +2697,184 @@
 t.MINUTES);%0A
+        %7D catch (AssertionError e) %7B%0A            throw new AssertionError(%22Failure in test setup: Failed to initialize at least 3 watcher nodes%22, e);%0A        %7D%0A
     %7D%0A%0A    @

Improve rolling upgrade test setup assertions (#58313)