Add better unit test to ensure task is scheduled/unscheduled.

Signed-off-by: Marc Handalian <handalm@amazon.com>
mch2 · Oct 19, 2023 · e4eaffc · e4eaffc
1 parent 22edb10
commit e4eaffc
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/server/src/test/java/org/opensearch/index/SegmentReplicationPressureServiceTests.java b/server/src/test/java/org/opensearch/index/SegmentReplicationPressureServiceTests.java
@@ -218,6 +218,7 @@ public void testFailStaleReplicaTask() throws Exception {
 
             // call the background task
             assertTrue(service.getFailStaleReplicaTask().mustReschedule());
+            assertTrue(service.getFailStaleReplicaTask().isScheduled());
             service.getFailStaleReplicaTask().runInternal();
 
             // verify that remote shard failed method is called which fails the replica shards falling behind.
@@ -281,14 +282,15 @@ public void testFailStaleReplicaTaskToggleOnOff() throws Exception {
             assertEquals(5, shardStats.getCheckpointsBehindCount());
 
             assertTrue(service.getFailStaleReplicaTask().mustReschedule());
-
+            assertTrue(service.getFailStaleReplicaTask().isScheduled());
             replicateSegments(primaryShard, shards.getReplicas());
 
             service.setReplicationTimeLimitFailReplica(TimeValue.ZERO);
             assertFalse(service.getFailStaleReplicaTask().mustReschedule());
-
+            assertFalse(service.getFailStaleReplicaTask().isScheduled());
             service.setReplicationTimeLimitFailReplica(TimeValue.timeValueMillis(1));
             assertTrue(service.getFailStaleReplicaTask().mustReschedule());
+            assertTrue(service.getFailStaleReplicaTask().isScheduled());
         }
     }