Fix integ tests

Signed-off-by: Sachin Kale <kalsac@amazon.com>
opensearch-project · Jul 20, 2023 · 8b6f741 · 8b6f741
1 parent 3a60440
commit 8b6f741
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 9 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreForceMergeIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreForceMergeIT.java
@@ -104,9 +104,17 @@ private void testRestoreWithMergeFlow(int numberOfIterations, boolean invokeFlus
         Map<String, Long> indexStats = indexData(numberOfIterations, invokeFlush, flushAfterMerge, deletedDocs);
 
         internalCluster().stopRandomNode(InternalTestCluster.nameFilter(primaryNodeName(INDEX_NAME)));
-        assertAcked(client().admin().indices().prepareClose(INDEX_NAME));
 
-        client().admin().cluster().restoreRemoteStore(new RestoreRemoteStoreRequest().indices(INDEX_NAME), PlainActionFuture.newFuture());
+        boolean restoreAllShards = randomBoolean();
+        if (restoreAllShards) {
+            assertAcked(client().admin().indices().prepareClose(INDEX_NAME));
+        }
+        client().admin()
+            .cluster()
+            .restoreRemoteStore(
+                new RestoreRemoteStoreRequest().indices(INDEX_NAME).restoreAllShards(restoreAllShards),
+                PlainActionFuture.newFuture()
+            );
         ensureGreen(INDEX_NAME);
 
         if (deletedDocs == -1) {

diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreIT.java
@@ -157,8 +157,16 @@ private void testRestoreFlow(boolean remoteTranslog, int numberOfIterations, boo
         internalCluster().stopRandomNode(InternalTestCluster.nameFilter(primaryNodeName(INDEX_NAME)));
         ensureRed(INDEX_NAME);
 
-        assertAcked(client().admin().indices().prepareClose(INDEX_NAME));
-        client().admin().cluster().restoreRemoteStore(new RestoreRemoteStoreRequest().indices(INDEX_NAME), PlainActionFuture.newFuture());
+        boolean restoreAllShards = randomBoolean();
+        if (restoreAllShards) {
+            assertAcked(client().admin().indices().prepareClose(INDEX_NAME));
+        }
+        client().admin()
+            .cluster()
+            .restoreRemoteStore(
+                new RestoreRemoteStoreRequest().indices(INDEX_NAME).restoreAllShards(restoreAllShards),
+                PlainActionFuture.newFuture()
+            );
 
         ensureGreen(INDEX_NAME);
         assertEquals(shardCount, getNumShards(INDEX_NAME).totalNumShards);

diff --git a/server/src/internalClusterTest/java/org/opensearch/snapshots/RestoreSnapshotIT.java b/server/src/internalClusterTest/java/org/opensearch/snapshots/RestoreSnapshotIT.java
@@ -281,7 +281,10 @@ public void testRestoreOperationsShallowCopyEnabled(boolean remoteTranslogEnable
         assertAcked(client.admin().indices().prepareClose(restoredIndexName1));
         client.admin()
             .cluster()
-            .restoreRemoteStore(new RestoreRemoteStoreRequest().indices(restoredIndexName1), PlainActionFuture.newFuture());
+            .restoreRemoteStore(
+                new RestoreRemoteStoreRequest().indices(restoredIndexName1).restoreAllShards(true),
+                PlainActionFuture.newFuture()
+            );
         ensureYellowAndNoInitializingShards(restoredIndexName1);
         ensureGreen(restoredIndexName1);
         assertDocsPresentInIndex(client(), restoredIndexName1, numDocsInIndex1);
@@ -443,7 +446,9 @@ public void testRestoreInSameRemoteStoreEnabledIndex() throws IOException {
         // Re-initialize client to make sure we are not using client from stopped node.
         client = client(clusterManagerNode);
         assertAcked(client.admin().indices().prepareClose(indexName1));
-        client.admin().cluster().restoreRemoteStore(new RestoreRemoteStoreRequest().indices(indexName1), PlainActionFuture.newFuture());
+        client.admin()
+            .cluster()
+            .restoreRemoteStore(new RestoreRemoteStoreRequest().indices(indexName1).restoreAllShards(true), PlainActionFuture.newFuture());
         ensureYellowAndNoInitializingShards(indexName1);
         ensureGreen(indexName1);
         assertDocsPresentInIndex(client(), indexName1, numDocsInIndex1);
@@ -524,7 +529,10 @@ public void testRestoreShallowCopySnapshotWithDifferentRepo() throws IOException
         assertAcked(client.admin().indices().prepareClose(restoredIndexName1));
         client.admin()
             .cluster()
-            .restoreRemoteStore(new RestoreRemoteStoreRequest().indices(restoredIndexName1), PlainActionFuture.newFuture());
+            .restoreRemoteStore(
+                new RestoreRemoteStoreRequest().indices(restoredIndexName1).restoreAllShards(true),
+                PlainActionFuture.newFuture()
+            );
         ensureYellowAndNoInitializingShards(restoredIndexName1);
         ensureGreen(restoredIndexName1);
         // indexing some new docs and validating

diff --git a/...va/org/opensearch/action/admin/cluster/remotestore/restore/RestoreRemoteStoreRequest.java b/...va/org/opensearch/action/admin/cluster/remotestore/restore/RestoreRemoteStoreRequest.java
@@ -35,8 +35,8 @@
 public class RestoreRemoteStoreRequest extends ClusterManagerNodeRequest<RestoreRemoteStoreRequest> implements ToXContentObject {
 
     private String[] indices = Strings.EMPTY_ARRAY;
-    private Boolean waitForCompletion;
-    private Boolean restoreAllShards;
+    private Boolean waitForCompletion = false;
+    private Boolean restoreAllShards = false;
 
     public RestoreRemoteStoreRequest() {}