Fix assertion in DiskThresholdDeciderIT (part 2) (elastic#127846) (elastic#128034)

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

Closes elastic#127787

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

Author: JeremyDahlgren

muted-tests.yml

@@ -452,9 +452,6 @@ tests:
 - class: org.elasticsearch.test.index.IndexVersionUtilsTests
   method: testIndexCompatibleVersionMatches
   issue: https://github.com/elastic/elasticsearch/issues/120760
-- class: org.elasticsearch.cluster.routing.allocation.decider.DiskThresholdDeciderIT
-  method: testRestoreSnapshotAllocationDoesNotExceedWatermarkWithMultipleRestores
-  issue: https://github.com/elastic/elasticsearch/issues/127787
 - class: org.elasticsearch.xpack.remotecluster.RemoteClusterSecurityFcActionAuthorizationIT
   method: testIndicesPrivilegesAreEnforcedForCcrRestoreSessionActions
   issue: https://github.com/elastic/elasticsearch/issues/127782

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) {