Fix tests

Signed-off-by: Rampreeth Ethiraj <ethirajrampreeth@gmail.com>
opensearch-project · Jul 16, 2024 · 2677ac5 · 2677ac5
1 parent 2481176
commit 2677ac5
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 14 deletions.
diff --git a/server/src/test/java/org/opensearch/index/seqno/PeerRecoveryRetentionLeaseExpiryTests.java b/server/src/test/java/org/opensearch/index/seqno/PeerRecoveryRetentionLeaseExpiryTests.java
@@ -151,7 +151,7 @@ public void testPeerRecoveryRetentionLeasesForAssignedCopiesDoNotEverExpire() {
         assertThat(
             leaseIds,
             equalTo(
-                replicationTracker.routingTable.shards()
+                replicationTracker.routingTable.assignedShards()
                     .stream()
                     .map(ReplicationTracker::getPeerRecoveryRetentionLeaseId)
                     .collect(Collectors.toSet())
@@ -183,7 +183,7 @@ public void testPeerRecoveryRetentionLeasesForUnassignedCopiesDoNotExpireImmedia
             equalTo(
                 Stream.concat(
                     Stream.of(ReplicationTracker.getPeerRecoveryRetentionLeaseId(unknownNodeId)),
-                    replicationTracker.routingTable.shards().stream().map(ReplicationTracker::getPeerRecoveryRetentionLeaseId)
+                    replicationTracker.routingTable.assignedShards().stream().map(ReplicationTracker::getPeerRecoveryRetentionLeaseId)
                 ).collect(Collectors.toSet())
             )
         );
@@ -215,7 +215,7 @@ public void testPeerRecoveryRetentionLeasesForUnassignedCopiesExpireEventually()
         assertThat(
             leaseIds,
             equalTo(
-                replicationTracker.routingTable.shards()
+                replicationTracker.routingTable.assignedShards()
                     .stream()
                     .map(ReplicationTracker::getPeerRecoveryRetentionLeaseId)
                     .collect(Collectors.toSet())
@@ -244,7 +244,7 @@ public void testPeerRecoveryRetentionLeasesForUnassignedCopiesExpireImmediatelyI
         assertThat(
             leaseIds,
             equalTo(
-                replicationTracker.routingTable.shards()
+                replicationTracker.routingTable.assignedShards()
                     .stream()
                     .map(ReplicationTracker::getPeerRecoveryRetentionLeaseId)
                     .collect(Collectors.toSet())
@@ -271,7 +271,7 @@ public void testPeerRecoveryRetentionLeasesForUnassignedCopiesExpireIfRetainingT
         assertThat(
             leaseIds,
             equalTo(
-                replicationTracker.routingTable.shards()
+                replicationTracker.routingTable.assignedShards()
                     .stream()
                     .map(ReplicationTracker::getPeerRecoveryRetentionLeaseId)
                     .collect(Collectors.toSet())

diff --git a/server/src/test/java/org/opensearch/index/seqno/ReplicationTrackerTestCase.java b/server/src/test/java/org/opensearch/index/seqno/ReplicationTrackerTestCase.java
@@ -104,13 +104,18 @@ static String nodeIdFromAllocationId(final AllocationId allocationId) {
     }
 
     static IndexShardRoutingTable routingTable(final Set<AllocationId> initializingIds, final AllocationId primaryId) {
-        return routingTable(initializingIds, Collections.singleton(primaryId), primaryId);
+        return routingTable(initializingIds, Collections.singleton(primaryId), primaryId, false);
+    }
+
+    static IndexShardRoutingTable routingTable(final Set<AllocationId> initializingIds, final AllocationId primaryId, final boolean shouldAddUnassignedShard) {
+        return routingTable(initializingIds, Collections.singleton(primaryId), primaryId, shouldAddUnassignedShard);
     }
 
     static IndexShardRoutingTable routingTable(
         final Set<AllocationId> initializingIds,
         final Set<AllocationId> activeIds,
-        final AllocationId primaryId
+        final AllocationId primaryId,
+        final boolean shouldAddUnassignedShard
     ) {
         final ShardId shardId = new ShardId("test", "_na_", 0);
         final ShardRouting primaryShard = TestShardRouting.newShardRouting(
@@ -121,13 +126,14 @@ static IndexShardRoutingTable routingTable(
             ShardRoutingState.STARTED,
             primaryId
         );
-        return routingTable(initializingIds, activeIds, primaryShard);
+        return routingTable(initializingIds, activeIds, primaryShard, shouldAddUnassignedShard);
     }
 
     static IndexShardRoutingTable routingTable(
         final Set<AllocationId> initializingIds,
         final Set<AllocationId> activeIds,
-        final ShardRouting primaryShard
+        final ShardRouting primaryShard,
+        final boolean shouldAddUnassignedShard
     ) {
         assert initializingIds != null && activeIds != null;
         assert !initializingIds.contains(primaryShard.allocationId());
@@ -136,7 +142,9 @@ static IndexShardRoutingTable routingTable(
         final IndexShardRoutingTable.Builder builder = new IndexShardRoutingTable.Builder(shardId);
 
         // Add a shard that is unassigned to simulate #11945
-        builder.addShard(TestShardRouting.newShardRoutingWithNullAllocationId(shardId, null, null, false, ShardRoutingState.UNASSIGNED));
+        if (shouldAddUnassignedShard) {
+            builder.addShard(TestShardRouting.newShardRoutingWithNullAllocationId(shardId, null, null, false, ShardRoutingState.UNASSIGNED));
+        }
 
         for (final AllocationId initializingId : initializingIds) {
             builder.addShard(

diff --git a/server/src/test/java/org/opensearch/index/seqno/ReplicationTrackerTests.java b/server/src/test/java/org/opensearch/index/seqno/ReplicationTrackerTests.java
@@ -988,7 +988,7 @@ private static FakeClusterState initialState() {
         return new FakeClusterState(
             initialClusterStateVersion,
             activeAllocationIds,
-            routingTable(initializingAllocationIds, Collections.singleton(primaryShard.allocationId()), primaryShard)
+            routingTable(initializingAllocationIds, Collections.singleton(primaryShard.allocationId()), primaryShard, true)
         );
     }
 
@@ -1031,7 +1031,8 @@ private static FakeClusterState randomUpdateClusterState(Set<String> allocationI
             routingTable(
                 Sets.difference(Sets.union(initializingIdsExceptRelocationTargets, initializingIdsToAdd), initializingIdsToRemove),
                 Collections.singleton(clusterState.routingTable.primaryShard().allocationId()),
-                clusterState.routingTable.primaryShard()
+                clusterState.routingTable.primaryShard(),
+                true
             )
         );
     }
@@ -1404,7 +1405,7 @@ public void testUpdateFromClusterManagerWithRemoteTranslogEnabled() {
             logger.info("  - [{}], local checkpoint [{}], [{}]", aId, allocations.get(aId), type);
         });
 
-        tracker.updateFromClusterManager(initialClusterStateVersion, ids(active), routingTable(initializing, active, primaryId));
+        tracker.updateFromClusterManager(initialClusterStateVersion, ids(active), routingTable(initializing, active, primaryId, true));
         tracker.activatePrimaryMode(NO_OPS_PERFORMED);
         assertEquals(tracker.getReplicationGroup().getReplicationTargets().size(), active.size());
         initializing.forEach(aId -> markAsTrackingAndInSyncQuietly(tracker, aId.getId(), NO_OPS_PERFORMED, false));
@@ -1615,7 +1616,7 @@ public void testInSyncIdsAreRemovedIfNotValidatedByClusterManagerWithRemoteTrans
             .put(IndexMetadata.SETTING_REMOTE_STORE_ENABLED, "true")
             .build();
         final ReplicationTracker tracker = newTracker(primaryId, settings, true);
-        tracker.updateFromClusterManager(initialClusterStateVersion, ids(active), routingTable(initializing, active, primaryId));
+        tracker.updateFromClusterManager(initialClusterStateVersion, ids(active), routingTable(initializing, active, primaryId, true));
         tracker.activatePrimaryMode(NO_OPS_PERFORMED);
         if (randomBoolean()) {
             initializingToStay.keySet().forEach(k -> markAsTrackingAndInSyncQuietly(tracker, k.getId(), NO_OPS_PERFORMED));