Fix flaky ShardIndexingPressureConcurrentExecutionTests (#5439) (#5450)

Add conditional check on assertNull to fix flaky tests. Signed-off-by: Rishikesh1159 <[email protected]> (cherry picked from commit 0bd3141) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: Rishikesh1159 <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Dec 5, 2022 · 1be8521 · 1be8521
1 parent f7e2d8e
commit 1be8521
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/server/src/test/java/org/opensearch/index/ShardIndexingPressureConcurrentExecutionTests.java b/server/src/test/java/org/opensearch/index/ShardIndexingPressureConcurrentExecutionTests.java
@@ -269,7 +269,13 @@ public void testCoordinatingPrimaryThreadedUpdateToShardLimitsAndRejections() th
 
         nodeStats = shardIndexingPressure.stats();
         IndexingPressurePerShardStats shardStoreStats = shardIndexingPressure.shardStats().getIndexingPressureShardStats(shardId1);
-        assertNull(shardStoreStats);
+        // If rejection count equals NUM_THREADS that means rejections happened until the last request, then we'll get shardStoreStats which
+        // was updated on the last request. In other cases, the shardStoreStats simply moves to the cold store and null is returned.
+        if (rejectionCount.get() == NUM_THREADS) {
+            assertEquals(10, shardStoreStats.getCurrentPrimaryAndCoordinatingLimits());
+        } else {
+            assertNull(shardStoreStats);
+        }
         shardStats = shardIndexingPressure.coldStats();
         if (randomBoolean) {
             assertEquals(rejectionCount.get(), nodeStats.getCoordinatingRejections());
@@ -331,7 +337,13 @@ public void testReplicaThreadedUpdateToShardLimitsAndRejections() throws Excepti
         assertEquals(0, nodeStats.getCurrentReplicaBytes());
 
         IndexingPressurePerShardStats shardStoreStats = shardIndexingPressure.shardStats().getIndexingPressureShardStats(shardId1);
-        assertNull(shardStoreStats);
+        // If rejection count equals NUM_THREADS that means rejections happened until the last request, then we'll get shardStoreStats which
+        // was updated on the last request. In other cases, the shardStoreStats simply moves to the cold store and null is returned.
+        if (rejectionCount.get() == NUM_THREADS) {
+            assertEquals(15, shardStoreStats.getCurrentReplicaLimits());
+        } else {
+            assertNull(shardStoreStats);
+        }
 
         shardStats = shardIndexingPressure.coldStats();
         assertEquals(rejectionCount.get(), shardStats.getIndexingPressureShardStats(shardId1).getReplicaNodeLimitsBreachedRejections());