From 918b71d0b196fe217a54159e6f147816acbc5d6f Mon Sep 17 00:00:00 2001
From: Suraj Singh <surajrider@gmail.com>
Date: Mon, 8 Jan 2024 11:45:07 -0800
Subject: [PATCH] Use SegmentReplicationTarget test to validate non-active
 on-disk files are reused for replication (#11786)

* Use SegmentReplicationTarget test to validate non-active on-disk files are reused for replication

Signed-off-by: Suraj Singh <surajrider@gmail.com>

* Fix original shard level unit test

Signed-off-by: Suraj Singh <surajrider@gmail.com>

---------

Signed-off-by: Suraj Singh <surajrider@gmail.com>
---
 .../index/shard/RemoteIndexShardTests.java    | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardTests.java b/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardTests.java
index dd92bfb47afdb..eacc504428ef1 100644
--- a/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardTests.java
+++ b/server/src/test/java/org/opensearch/index/shard/RemoteIndexShardTests.java
@@ -11,6 +11,7 @@
 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.index.SegmentInfos;
 import org.apache.lucene.util.Version;
+import org.opensearch.action.StepListener;
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.concurrent.GatedCloseable;
 import org.opensearch.common.settings.Settings;
@@ -371,7 +372,6 @@ public void testPrimaryRestart() throws Exception {
      * prevent FileAlreadyExistsException. It does so by only copying files in first round of segment replication without
      * committing locally so that in next round of segment replication those files are not considered for download again
      */
-    @AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/10885")
     public void testSegRepSucceedsOnPreviousCopiedFiles() throws Exception {
         try (ReplicationGroup shards = createGroup(1, getIndexSettings(), new NRTReplicationEngineFactory())) {
             shards.startAll();
@@ -388,6 +388,7 @@ public void testSegRepSucceedsOnPreviousCopiedFiles() throws Exception {
             );
             CountDownLatch latch = new CountDownLatch(1);
 
+            logger.info("--> Starting first round of replication");
             // Start first round of segment replication. This should fail with simulated error but with replica having
             // files in its local store but not in active reader.
             final SegmentReplicationTarget target = targetService.startReplication(
@@ -427,6 +428,7 @@ public void onReplicationFailure(
             // Start next round of segment replication and not throwing exception resulting in commit on replica
             when(sourceFactory.get(any())).thenReturn(getRemoteStoreReplicationSource(replica, () -> {}));
             CountDownLatch waitForSecondRound = new CountDownLatch(1);
+            logger.info("--> Starting second round of replication");
             final SegmentReplicationTarget newTarget = targetService.startReplication(
                 replica,
                 primary.getLatestReplicationCheckpoint(),
@@ -560,8 +562,19 @@ public void getSegmentFiles(
                 BiConsumer<String, Long> fileProgressTracker,
                 ActionListener<GetSegmentFilesResponse> listener
             ) {
-                super.getSegmentFiles(replicationId, checkpoint, filesToFetch, indexShard, (fileName, bytesRecovered) -> {}, listener);
-                postGetFilesRunnable.run();
+                StepListener<GetSegmentFilesResponse> waitForCopyFilesListener = new StepListener();
+                super.getSegmentFiles(
+                    replicationId,
+                    checkpoint,
+                    filesToFetch,
+                    indexShard,
+                    (fileName, bytesRecovered) -> {},
+                    waitForCopyFilesListener
+                );
+                waitForCopyFilesListener.whenComplete(response -> {
+                    postGetFilesRunnable.run();
+                    listener.onResponse(response);
+                }, listener::onFailure);
             }
 
             @Override