diff --git a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreRestoreIT.java b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreRestoreIT.java index 65335f444a2df..c5143f9755a4d 100644 --- a/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreRestoreIT.java +++ b/server/src/internalClusterTest/java/org/opensearch/remotestore/RemoteStoreRestoreIT.java @@ -13,6 +13,7 @@ import org.opensearch.action.support.PlainActionFuture; import org.opensearch.cluster.health.ClusterHealthStatus; import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.action.ActionListener; import org.opensearch.repositories.RepositoriesService; import org.opensearch.test.InternalTestCluster; import org.opensearch.test.OpenSearchIntegTestCase; @@ -30,6 +31,30 @@ @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 0) public class RemoteStoreRestoreIT extends BaseRemoteStoreRestoreIT { + public void testWaitForCompletion() throws IOException, ExecutionException, InterruptedException { + prepareCluster(1, 3, INDEX_NAME, 0, 1); + Map indexStats = indexData(2, true, INDEX_NAME); + + assertHitCount(client().prepareSearch(INDEX_NAME).setSize(0).get(), indexStats.get(REFRESHED_OR_FLUSHED_OPERATIONS)); + + internalCluster().stopRandomNode(InternalTestCluster.nameFilter(primaryNodeName(INDEX_NAME))); + ensureRed(INDEX_NAME); + + PlainActionFuture listener = PlainActionFuture.newFuture(); + client().admin() + .cluster() + .restoreRemoteStore( + new RestoreRemoteStoreRequest().indices(INDEX_NAME).waitForCompletion(true), + listener + ); + RestoreRemoteStoreResponse restoreRemoteStoreResponse = listener.get(); + + assertEquals(1, restoreRemoteStoreResponse.getRestoreInfo().successfulShards()); + assertEquals(0, restoreRemoteStoreResponse.getRestoreInfo().failedShards()); + assertHitCount(client().prepareSearch(INDEX_NAME).setSize(0).get(), indexStats.get(TOTAL_OPERATIONS)); + logger.info("--> Successful"); + } + /** * Simulates all data restored using Remote Translog Store. * @throws IOException IO Exception. diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java index 2b0e5c74cad53..aaa7aad9a34ac 100644 --- a/server/src/main/java/org/opensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java +++ b/server/src/main/java/org/opensearch/action/admin/cluster/remotestore/restore/TransportRestoreRemoteStoreAction.java @@ -85,7 +85,7 @@ protected void clusterManagerOperation( final ActionListener listener ) { restoreService.restore(request, ActionListener.delegateFailure(listener, (delegatedListener, restoreCompletionResponse) -> { - if (restoreCompletionResponse.getRestoreInfo() == null && request.waitForCompletion()) { + if (request.waitForCompletion()) { RestoreClusterStateListener.createAndRegisterListener( clusterService, restoreCompletionResponse,