test fixes

Signed-off-by: Ayush Kataria <[email protected]>
opensearch-project · Oct 10, 2022 · a06998e · a06998e
1 parent 2d27461
commit a06998e
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 8 deletions.
diff --git a/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java b/server/src/main/java/org/opensearch/indices/replication/SegmentReplicationTargetService.java
@@ -12,7 +12,6 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.opensearch.ExceptionsHelper;
-import org.opensearch.OpenSearchException;
 import org.opensearch.action.ActionListener;
 import org.opensearch.cluster.routing.ShardRouting;
 import org.opensearch.common.Nullable;
@@ -197,7 +196,7 @@ public void onReplicationDone(SegmentReplicationState state) {
                 }
 
                 @Override
-                public void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure) {
+                public void onReplicationFailure(SegmentReplicationState state, ReplicationFailedException e, boolean sendShardFailure) {
                     logger.trace(
                         () -> new ParameterizedMessage(
                             "[shardId {}] [replication id {}] Replication failed, timing data: {}",
@@ -256,7 +255,7 @@ default void onFailure(ReplicationState state, ReplicationFailedException e, boo
 
         void onReplicationDone(SegmentReplicationState state);
 
-        void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure);
+        void onReplicationFailure(SegmentReplicationState state, ReplicationFailedException e, boolean sendShardFailure);
     }
 
     /**

diff --git a/server/src/test/java/org/opensearch/index/shard/SegmentReplicationIndexShardTests.java b/server/src/test/java/org/opensearch/index/shard/SegmentReplicationIndexShardTests.java
@@ -44,6 +44,7 @@
 import org.opensearch.indices.replication.checkpoint.SegmentReplicationCheckpointPublisher;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
 import org.opensearch.indices.replication.common.CopyState;
+import org.opensearch.indices.replication.common.ReplicationFailedException;
 import org.opensearch.indices.replication.common.ReplicationType;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
@@ -670,8 +671,7 @@ public void onReplicationDone(SegmentReplicationState state) {
                 }
 
                 @Override
-                public void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure) {
-                    assertTrue(e instanceof CancellableThreads.ExecutionCancelledException);
+                public void onReplicationFailure(SegmentReplicationState state, ReplicationFailedException e, boolean sendShardFailure) {
                     assertFalse(sendShardFailure);
                     assertEquals(SegmentReplicationState.Stage.CANCELLED, state.getStage());
                     latch.countDown();

diff --git a/...rc/test/java/org/opensearch/indices/replication/SegmentReplicationTargetServiceTests.java b/...rc/test/java/org/opensearch/indices/replication/SegmentReplicationTargetServiceTests.java
@@ -22,6 +22,7 @@
 import org.opensearch.index.store.Store;
 import org.opensearch.index.store.StoreFileMetadata;
 import org.opensearch.indices.replication.checkpoint.ReplicationCheckpoint;
+import org.opensearch.indices.replication.common.ReplicationFailedException;
 import org.opensearch.indices.replication.common.ReplicationType;
 
 import java.io.IOException;
@@ -104,7 +105,7 @@ public void onReplicationDone(SegmentReplicationState state) {
             }
 
             @Override
-            public void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure) {
+            public void onReplicationFailure(SegmentReplicationState state, ReplicationFailedException e, boolean sendShardFailure) {
                 logger.error("Unexpected error", e);
                 Assert.fail("Test should succeed");
             }
@@ -149,7 +150,7 @@ public void onReplicationDone(SegmentReplicationState state) {
                 }
 
                 @Override
-                public void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure) {
+                public void onReplicationFailure(SegmentReplicationState state, ReplicationFailedException e, boolean sendShardFailure) {
                     // failures leave state object in last entered stage.
                     assertEquals(SegmentReplicationState.Stage.GET_CHECKPOINT_INFO, state.getStage());
                     assertEquals(expectedError, e.getCause());

diff --git a/test/framework/src/main/java/org/opensearch/index/shard/IndexShardTestCase.java b/test/framework/src/main/java/org/opensearch/index/shard/IndexShardTestCase.java
@@ -1283,7 +1283,7 @@ public void onReplicationDone(SegmentReplicationState state) {
                     }
 
                     @Override
-                    public void onReplicationFailure(SegmentReplicationState state, OpenSearchException e, boolean sendShardFailure) {
+                    public void onReplicationFailure(SegmentReplicationState state, ReplicationFailedException e, boolean sendShardFailure) {
                         logger.error("Unexpected replication failure in test", e);
                         Assert.fail("test replication should not fail: " + e);
                     }