Fix assertion in DiskThresholdDeciderIT (part 2) (#127846) (#128033)

Fixes the final assertion in
testRestoreSnapshotAllocationDoesNotExceedWatermarkWithMultipleRestores()
that wasn't addressed with the fix in PR 127615 for issue 127286.

Closes #127787

Co-authored-by: David Turner <[email protected]>
(cherry picked from commit 98296d7)

Author: JeremyDahlgren

server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderIT.java

@@ -280,7 +280,12 @@ public void testRestoreSnapshotAllocationDoesNotExceedWatermarkWithMultipleResto
             tinyNodeShardIds.size() + tinyNodeShardIdsCopy.size(),
             is(1)
         );
-        assertThat(tinyNodeShardIds.iterator().next(), in(shardSizes.getShardIdsWithSizeSmallerOrEqual(usableSpace)));
+        final var useableSpaceShardSizes = shardSizes.getShardIdsWithSizeSmallerOrEqual(usableSpace);
+        final var tinyNodeShardId = tinyNodeShardIds.isEmpty() == false
+            ? tinyNodeShardIds.iterator().next()
+            // shardSizes only contains the sizes from the original index, not the copy, so we map the copied shard back to the original idx
+            : new ShardId(useableSpaceShardSizes.iterator().next().getIndex(), tinyNodeShardIdsCopy.iterator().next().id());
+        assertThat(tinyNodeShardId, in(useableSpaceShardSizes));
     }
 
     private Set<ShardId> getShardIds(final String nodeId, final String indexName) {