Incorporated comments

Signed-off-by: Dhwanil Patel <dhwanip@amazon.com>

server/src/internalClusterTest/java/org/opensearch/action/admin/indices/datastream/DataStreamTestCase.java

@@ -30,23 +30,14 @@
 import static org.opensearch.test.OpenSearchIntegTestCase.Scope;
 import static org.hamcrest.Matchers.is;
 
-@ClusterScope(scope = Scope.TEST, numDataNodes = DataStreamTestCase.DATA_NODE_COUNT)
+@ClusterScope(scope = Scope.TEST, numDataNodes = 2)
 public class DataStreamTestCase extends OpenSearchIntegTestCase {
 
- protected static final int DATA_NODE_COUNT = 2;
-
- private void resetCluster() {
- internalCluster().stopAllNodes();
- internalCluster().startNodes(DATA_NODE_COUNT);
- }
-
  public AcknowledgedResponse createDataStream(String name) throws Exception {
  CreateDataStreamAction.Request request = new CreateDataStreamAction.Request(name);
  AcknowledgedResponse response = client().admin().indices().createDataStream(request).get();
  assertThat(response.isAcknowledged(), is(true));
- if (performRemoteStateRestore()) {
- resetCluster();
- }
+ performRemoteStoreTestAction();
  return response;
  }
 
@@ -77,12 +68,7 @@ public RolloverResponse rolloverDataStream(String name) throws Exception {
  RolloverResponse response = client().admin().indices().rolloverIndex(request).get();
  assertThat(response.isAcknowledged(), is(true));
  assertThat(response.isRolledOver(), is(true));
- if (performRemoteStateRestore()) {
- String clusterUUIDBefore = clusterService().state().metadata().clusterUUID();
- resetCluster();
- String clusterUUIDAfter = clusterService().state().metadata().clusterUUID();
- assertFalse(clusterUUIDBefore.equals(clusterUUIDAfter));
- }
+ performRemoteStoreTestAction();
  return response;
  }
 
@@ -125,8 +111,4 @@ public AcknowledgedResponse deleteIndexTemplate(String name) throws Exception {
  assertThat(response.isAcknowledged(), is(true));
  return response;
  }
-
- protected boolean performRemoteStateRestore() {
- return false;
- }
 }

server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreClusterStateRestoreIT.java

@@ -216,7 +216,7 @@ private void validateCurrentMetadata() throws Exception {
  });
  }
 
- public void testDataStreamWithRemoteStateRestore() throws Exception {
+ public void testDataStreamPostRemoteStateRestore() throws Exception {
  new DataStreamRolloverIT() {
  protected boolean performRemoteStateRestore() {
  return true;

test/framework/src/main/java/org/opensearch/test/InternalTestCluster.java

@@ -1871,6 +1871,18 @@ public void stopAllNodes() {
  }
  }
 
+ /**
+ * Replace all nodes by stopping all current node and starting new node.
+ * Used for remote store test cases, where remote state is restored.
+ */
+ public void resetCluster() {
+ int totalClusterManagerNodes = numClusterManagerNodes();
+ int totalDataNodes = numDataNodes();
+ stopAllNodes();
+ startClusterManagerOnlyNodes(totalClusterManagerNodes);
+ startDataOnlyNodes(totalDataNodes);
+ }
+
  private synchronized void startAndPublishNodesAndClients(List<NodeAndClient> nodeAndClients) {
  if (nodeAndClients.size() > 0) {
  final int newClusterManagers = (int) nodeAndClients.stream()

test/framework/src/main/java/org/opensearch/test/OpenSearchIntegTestCase.java

@@ -789,6 +789,30 @@ protected Settings featureFlagSettings() {
  return featureSettings.build();
  }
 
+ /**
+ * Represent if it needs to trigger remote state restore or not.
+ * For tests with remote store enabled domain, it will be overridden to true.
+ *
+ * @return if needs to perform remote state restore or not
+ */
+ protected boolean triggerRemoteStateRestore() {
+ return false;
+ }
+
+ /**
+ * For tests with remote cluster state, it will reset the cluster and cluster state will be
+ * restored from remote.
+ */
+ protected void performRemoteStoreTestAction() {
+ if(triggerRemoteStateRestore()) {
+ String clusterUUIDBefore = clusterService().state().metadata().clusterUUID();
+ internalCluster().resetCluster();
+ String clusterUUIDAfter = clusterService().state().metadata().clusterUUID();
+ // assertion that UUID is changed post restore.
+ assertFalse(clusterUUIDBefore.equals(clusterUUIDAfter));
+ }
+ }
+
  /**
  * Creates one or more indices and asserts that the indices are acknowledged. If one of the indices
  * already exists this method will fail and wipe all the indices created so far.