Fix testDesiredBalanceShouldConvergeInABigCluster (#104442)

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputerTests.java

@@ -581,7 +581,6 @@ public class DesiredBalanceComputerTests extends ESAllocationTestCase {
         );
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/104343")
     public void testDesiredBalanceShouldConvergeInABigCluster() {
         var nodes = randomIntBetween(3, 7);
         var nodeIds = new ArrayList<String>(nodes);
@@ -639,18 +638,46 @@ public class DesiredBalanceComputerTests extends ESAllocationTestCase {
                 if (primaryNodeId != null) {
                     dataPath.put(new NodeAndShard(primaryNodeId, shardId), "/data");
                     usedDiskSpace.compute(primaryNodeId, (k, v) -> v + thisShardSize);
+                    indexRoutingTableBuilder.addShard(
+                        newShardRouting(
+                            shardId,
+                            primaryNodeId,
+                            null,
+                            true,
+                            STARTED,
+                            AllocationId.newInitializing(inSyncIds.get(shard * (replicas + 1)))
+                        )
+                    );
+                } else {
+                    var lastAllocatedNodeId = randomFrom(remainingNodeIds);
+                    assertThat(lastAllocatedNodeId, notNullValue());// the only null was picked as primaryNodeId
+                    dataPath.put(new NodeAndShard(lastAllocatedNodeId, shardId), "/data");
+                    usedDiskSpace.compute(lastAllocatedNodeId, (k, v) -> v + thisShardSize);
+                    indexRoutingTableBuilder.addShard(
+                        newShardRouting(
+                            shardId,
+                            null,
+                            null,
+                            true,
+                            UNASSIGNED,
+                            RecoverySource.ExistingStoreRecoverySource.INSTANCE,
+                            new UnassignedInfo(
+                                UnassignedInfo.Reason.NODE_LEFT,
+                                null,
+                                null,
+                                0,
+                                0,
+                                0,
+                                false,
+                                UnassignedInfo.AllocationStatus.NO_ATTEMPT,
+                                Set.of(),
+                                lastAllocatedNodeId
+                            ),
+                            AllocationId.newInitializing(inSyncIds.get(shard * (replicas + 1)))
+                        )
+                    );
                 }
 
-                indexRoutingTableBuilder.addShard(
-                    newShardRouting(
-                        shardId,
-                        primaryNodeId,
-                        null,
-                        true,
-                        primaryNodeId == null ? UNASSIGNED : STARTED,
-                        AllocationId.newInitializing(inSyncIds.get(shard * (replicas + 1)))
-                    )
-                );
                 for (int replica = 0; replica < replicas; replica++) {
                     var replicaNodeId = primaryNodeId == null ? null : pickAndRemoveRandomValueFrom(remainingNodeIds);
                     shardSizes.put(shardIdentifierFromRouting(shardId, false), thisShardSize);

test/framework/src/main/java/org/elasticsearch/cluster/routing/TestShardRouting.java

@@ -287,6 +287,31 @@ public class TestShardRouting {
         );
     }
 
+    public static ShardRouting newShardRouting(
+        ShardId shardId,
+        String currentNodeId,
+        String relocatingNodeId,
+        boolean primary,
+        ShardRoutingState state,
+        RecoverySource recoverySource,
+        UnassignedInfo unassignedInfo,
+        AllocationId allocationId
+    ) {
+        return new ShardRouting(
+            shardId,
+            currentNodeId,
+            relocatingNodeId,
+            primary,
+            state,
+            recoverySource,
+            unassignedInfo,
+            buildRelocationFailureInfo(state),
+            allocationId,
+            -1,
+            ShardRouting.Role.DEFAULT
+        );
+    }
+
     public static ShardRouting relocate(ShardRouting shardRouting, String relocatingNodeId, long expectedShardSize) {
         return shardRouting.relocate(relocatingNodeId, expectedShardSize);
     }