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

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]>

Author: JeremyDahlgren

muted-tests.yml

@@ -447,9 +447,6 @@ tests:
 - class: org.elasticsearch.xpack.esql.qa.mixed.EsqlClientYamlIT
   method: test {p0=esql/120_profile/avg 8.14 or after}
   issue: https://github.com/elastic/elasticsearch/issues/127879
-- class: org.elasticsearch.cluster.routing.allocation.decider.DiskThresholdDeciderIT
-  method: testRestoreSnapshotAllocationDoesNotExceedWatermarkWithMultipleRestores
-  issue: https://github.com/elastic/elasticsearch/issues/127787
 - class: org.elasticsearch.indices.stats.IndexStatsIT
   method: testThrottleStats
   issue: https://github.com/elastic/elasticsearch/issues/126359

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