[Segment Replication] Wait for segment replication to be completed an…

…d marked done before assertion (#6370) (#6373)

* [Segment Replication] Wait for segment replication to be completed and marked done before assertion



* Spotless fix



---------

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

server/src/internalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java

@@ -18,10 +18,7 @@
 import org.opensearch.transport.TransportService;
 
 import java.util.concurrent.CountDownLatch;
-
 import static java.util.Arrays.asList;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
 
 @OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0)
 public class SegmentReplicationStatsIT extends SegmentReplicationBaseIT {
@@ -40,20 +37,17 @@ public void testSegmentReplicationStatsResponse() throws Exception {
         refresh(INDEX_NAME);
         waitForSearchableDocs(10L, asList(primaryNode, replicaNode));
 
-        SegmentReplicationStatsResponse response = client().admin()
-            .indices()
-            .prepareSegmentReplicationStats(INDEX_NAME)
-            .execute()
-            .actionGet();
-        // Verify API Response
-        assertThat(response.shardSegmentReplicationStates().size(), equalTo(SHARD_COUNT));
-        assertBusy(
-            () -> assertThat(
-                response.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getStage(),
-                equalTo(SegmentReplicationState.Stage.DONE)
-            )
-        );
-        assertThat(response.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getIndex().recoveredFileCount(), greaterThan(0));
+        assertBusy(() -> {
+            final SegmentReplicationStatsResponse response = client().admin()
+                .indices()
+                .prepareSegmentReplicationStats(INDEX_NAME)
+                .execute()
+                .actionGet();
+            // Verify API Response
+            assertEquals(response.shardSegmentReplicationStates().size(), SHARD_COUNT);
+            assertEquals(response.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getStage(), SegmentReplicationState.Stage.DONE);
+            assertTrue(response.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getIndex().recoveredFileCount() > 0);
+        });
     }
 
     public void testSegmentReplicationStatsResponseForActiveAndCompletedOnly() throws Exception {
@@ -110,9 +104,9 @@ public void testSegmentReplicationStatsResponseForActiveAndCompletedOnly() throw
             .setActiveOnly(true)
             .execute()
             .actionGet();
-        assertThat(
+        assertEquals(
             activeOnlyResponse.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getStage(),
-            equalTo(SegmentReplicationState.Stage.GET_FILES)
+            SegmentReplicationState.Stage.GET_FILES
         );
 
         // verifying completed_only by checking if current stage is DONE
@@ -122,10 +116,10 @@ public void testSegmentReplicationStatsResponseForActiveAndCompletedOnly() throw
             .setCompletedOnly(true)
             .execute()
             .actionGet();
-        assertThat(completedOnlyResponse.shardSegmentReplicationStates().size(), equalTo(SHARD_COUNT));
-        assertThat(
+        assertEquals(completedOnlyResponse.shardSegmentReplicationStates().size(), SHARD_COUNT);
+        assertEquals(
             completedOnlyResponse.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getStage(),
-            equalTo(SegmentReplicationState.Stage.DONE)
+            SegmentReplicationState.Stage.DONE
         );
         waitForAssertions.countDown();
     }