From 7216b78089f1be17536ecc04a4a5f3bd9405e1c0 Mon Sep 17 00:00:00 2001
From: Vacha <vachshah@amazon.com>
Date: Tue, 10 Aug 2021 14:09:24 -0700
Subject: [PATCH] Adding support for restart upgrades for plugins bwc

Signed-off-by: Vacha <vachshah@amazon.com>
---
 .../testclusters/OpenSearchCluster.java       | 38 +++++++++++--------
 1 file changed, 22 insertions(+), 16 deletions(-)

diff --git a/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchCluster.java b/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchCluster.java
index 0c46769e10aed..1e3d714fe186b 100644
--- a/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchCluster.java
+++ b/buildSrc/src/main/java/org/opensearch/gradle/testclusters/OpenSearchCluster.java
@@ -392,32 +392,26 @@ public void goToNextVersion() {
         writeUnicastHostsFiles();
     }
 
+    public void upgradeAllNodesAndPluginsToNextVersion(List<Provider<RegularFile>> plugins) {
+        stop(false);
+        nodes.all(OpenSearchNode::goToNextVersion);
+        upgradePlugin(plugins);
+        start();
+        writeUnicastHostsFiles();
+    }
+
     public void fullRestart() {
         stop(false);
         start();
     }
 
     public void nextNodeToNextVersion() {
-        if (nodeIndex + 1 > nodes.size()) {
-            throw new TestClustersException("Ran out of nodes to take to the next version");
-        }
-        OpenSearchNode node = nodes.getByName(clusterName + "-" + nodeIndex);
-        node.stop(false);
-        node.goToNextVersion();
-        commonNodeConfig(node, null, null);
-        nodeIndex += 1;
+        OpenSearchNode node = upgradeNodeToNextVersion();
         node.start();
     }
 
     public void upgradeNodeAndPluginToNextVersion(List<Provider<RegularFile>> plugins) {
-        if (nodeIndex + 1 > nodes.size()) {
-            throw new TestClustersException("Ran out of nodes to take to the next version");
-        }
-        OpenSearchNode node = nodes.getByName(clusterName + "-" + nodeIndex);
-        node.stop(false);
-        node.goToNextVersion();
-        commonNodeConfig(node, null, null);
-        nodeIndex += 1;
+        OpenSearchNode node = upgradeNodeToNextVersion();
         node.upgradePlugin(plugins);
         node.start();
     }
@@ -453,6 +447,18 @@ private void writeUnicastHostsFiles() {
         });
     }
 
+    private OpenSearchNode upgradeNodeToNextVersion() {
+        if (nodeIndex + 1 > nodes.size()) {
+            throw new TestClustersException("Ran out of nodes to take to the next version");
+        }
+        OpenSearchNode node = nodes.getByName(clusterName + "-" + nodeIndex);
+        node.stop(false);
+        node.goToNextVersion();
+        commonNodeConfig(node, null, null);
+        nodeIndex += 1;
+        return node;
+    }
+
     @Override
     @Internal
     public String getHttpSocketURI() {