Fix flaky Tests.

Signed-off-by: Rishikesh1159 <[email protected]>
opensearch-project · Feb 9, 2023 · ad55487 · ad55487
1 parent 84a8d3e
commit ad55487
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git a/server/src/internalClusterTest/java/org/opensearch/action/admin/ClientTimeoutIT.java b/server/src/internalClusterTest/java/org/opensearch/action/admin/ClientTimeoutIT.java
@@ -51,6 +51,11 @@ protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Collections.singletonList(MockTransportService.TestPlugin.class);
     }
 
+    @Override
+    protected boolean addMockInternalEngine() {
+        return false;
+    }
+
     public void testNodesInfoTimeout() {
         String clusterManagerNode = internalCluster().startClusterManagerOnlyNode();
         String dataNode = internalCluster().startDataOnlyNode();

diff --git a/...nternalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java b/...nternalClusterTest/java/org/opensearch/indices/replication/SegmentReplicationStatsIT.java
@@ -47,7 +47,12 @@ public void testSegmentReplicationStatsResponse() throws Exception {
             .actionGet();
         // Verify API Response
         assertThat(response.shardSegmentReplicationStates().size(), equalTo(SHARD_COUNT));
-        assertThat(response.shardSegmentReplicationStates().get(INDEX_NAME).get(0).getStage(), equalTo(SegmentReplicationState.Stage.DONE));
+        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));
     }
 

diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationState.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationState.java
@@ -22,8 +22,8 @@
 import org.opensearch.indices.replication.common.ReplicationTimer;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * ReplicationState implementation to track Segment Replication events.
@@ -76,9 +76,10 @@ public static Stage fromId(byte id) {
 
     private Stage stage;
     private ReplicationLuceneIndex index;
-
     private final ReplicationTimer overallTimer;
-    private final Map<String, Long> timingData;
+
+    // Timing data will have as many entries as stages, plus one
+    private final Map<String, Long> timingData = new ConcurrentHashMap<>(Stage.values().length + 1);
     private final ReplicationTimer stageTimer;
     private long replicationId;
     private final ShardRouting shardRouting;
@@ -152,8 +153,6 @@ public SegmentReplicationState(
         this.replicationId = replicationId;
         this.sourceDescription = sourceDescription;
         this.targetNode = targetNode;
-        // Timing data will have as many entries as stages, plus one
-        timingData = new HashMap<>(Stage.values().length + 1);
         overallTimer = new ReplicationTimer();
         stageTimer = new ReplicationTimer();
         setStage(Stage.INIT);
@@ -167,7 +166,6 @@ public SegmentReplicationState(StreamInput in) throws IOException {
         replicationId = in.readLong();
         overallTimer = new ReplicationTimer(in);
         stageTimer = new ReplicationTimer(in);
-        timingData = in.readMap(StreamInput::readString, StreamInput::readLong);
         sourceDescription = in.readString();
         targetNode = new DiscoveryNode(in);
     }
@@ -180,7 +178,6 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeLong(replicationId);
         overallTimer.writeTo(out);
         stageTimer.writeTo(out);
-        out.writeMap(timingData, StreamOutput::writeString, StreamOutput::writeLong);
         out.writeString(sourceDescription);
         targetNode.writeTo(out);
     }