Fix test

Signed-off-by: Anshu Agarwal <[email protected]>
opensearch-project · Nov 10, 2022 · c0e5b67 · c0e5b67
1 parent 5ee2b02
commit c0e5b67
Showing 1 changed file with 19 additions and 24 deletions.
diff --git a/server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java b/server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java
@@ -647,13 +647,12 @@ public void testWeightedRoutingInMemoryStore() {
                 .index("test")
                 .shard(0)
                 .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1);
-            assertEquals(2, shardIterator.size());
+            assertEquals(3, shardIterator.size());
             ShardRouting shardRouting;
-            while (shardIterator.remaining() > 0) {
-                shardRouting = shardIterator.nextOrNull();
-                assertNotNull(shardRouting);
-                assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId()));
-            }
+            shardRouting = shardIterator.nextOrNull();
+            assertNotNull(shardRouting);
+            assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId()));
+
 
             // Make iterator call with same WeightedRouting instance
             assertNotNull(
@@ -663,12 +662,10 @@ public void testWeightedRoutingInMemoryStore() {
                 .index("test")
                 .shard(0)
                 .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1);
-            assertEquals(2, shardIterator.size());
-            while (shardIterator.remaining() > 0) {
-                shardRouting = shardIterator.nextOrNull();
-                assertNotNull(shardRouting);
-                assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId()));
-            }
+            assertEquals(3, shardIterator.size());
+            shardRouting = shardIterator.nextOrNull();
+            assertNotNull(shardRouting);
+            assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId()));
 
             // Make iterator call with new instance of WeightedRouting but same weights
             Map<String, Double> weights1 = Map.of("zone1", 1.0, "zone2", 1.0, "zone3", 0.0);
@@ -680,12 +677,11 @@ public void testWeightedRoutingInMemoryStore() {
                 .index("test")
                 .shard(0)
                 .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1);
-            assertEquals(2, shardIterator.size());
-            while (shardIterator.remaining() > 0) {
-                shardRouting = shardIterator.nextOrNull();
-                assertNotNull(shardRouting);
-                assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId()));
-            }
+            assertEquals(3, shardIterator.size());
+            shardRouting = shardIterator.nextOrNull();
+            assertNotNull(shardRouting);
+            assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId()));
+
 
             // Make iterator call with different weights
             Map<String, Double> weights2 = Map.of("zone1", 1.0, "zone2", 0.0, "zone3", 1.0);
@@ -697,12 +693,11 @@ public void testWeightedRoutingInMemoryStore() {
                 .index("test")
                 .shard(0)
                 .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1);
-            assertEquals(2, shardIterator.size());
-            while (shardIterator.remaining() > 0) {
-                shardRouting = shardIterator.nextOrNull();
-                assertNotNull(shardRouting);
-                assertFalse(Arrays.asList("node2").contains(shardRouting.currentNodeId()));
-            }
+            assertEquals(3, shardIterator.size());
+            shardRouting = shardIterator.nextOrNull();
+            assertNotNull(shardRouting);
+            assertFalse(Arrays.asList("node2").contains(shardRouting.currentNodeId()));
+
         } finally {
             terminate(threadPool);
         }