[Test] Allow allocation in mixed cluster

muted-tests.yml

@@ -571,15 +571,6 @@ tests:
 - class: org.elasticsearch.server.cli.MachineDependentHeapTests
   method: testMlOnlyOptions
   issue: https://github.com/elastic/elasticsearch/issues/129236
-- class: org.elasticsearch.upgrades.RunningSnapshotIT
-  method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=1}
-  issue: https://github.com/elastic/elasticsearch/issues/129644
-- class: org.elasticsearch.upgrades.RunningSnapshotIT
-  method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=2}
-  issue: https://github.com/elastic/elasticsearch/issues/129645
-- class: org.elasticsearch.upgrades.RunningSnapshotIT
-  method: testRunningSnapshotCompleteAfterUpgrade {upgradedNodes=3}
-  issue: https://github.com/elastic/elasticsearch/issues/129646
 - class: org.elasticsearch.test.apmintegration.TracesApmIT
   method: testApmIntegration
   issue: https://github.com/elastic/elasticsearch/issues/129651

qa/rolling-upgrade/src/javaRestTest/java/org/elasticsearch/upgrades/RunningSnapshotIT.java

@@ -18,10 +18,15 @@
 
 import java.io.IOException;
 import java.util.Collection;
+import java.util.Map;
+import java.util.stream.Collectors;
 
+import static org.elasticsearch.client.RestClient.IGNORE_RESPONSE_CODES_PARAM;
 import static org.elasticsearch.upgrades.SnapshotBasedRecoveryIT.indexDocs;
+import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
 import static org.hamcrest.Matchers.not;
 
 public class RunningSnapshotIT extends AbstractRollingUpgradeTestCase {
@@ -34,7 +39,11 @@ public void testRunningSnapshotCompleteAfterUpgrade() throws Exception {
         final String indexName = "index";
         final String repositoryName = "repo";
         final String snapshotName = "snapshot";
-        final var nodeIds = getNodesInfo(client()).keySet();
+        final Map<String, Map<?, ?>> nodesInfo = getNodesInfo(client());
+        final var nodeIdToNodeNames = nodesInfo.entrySet()
+            .stream()
+            .collect(Collectors.toUnmodifiableMap(Map.Entry::getKey, entry -> entry.getValue().get("name").toString()));
+        assertThat(nodeIdToNodeNames.values(), containsInAnyOrder("test-cluster-0", "test-cluster-1", "test-cluster-2"));
 
         if (isOldCluster()) {
             registerRepository(repositoryName, "fs", randomBoolean(), Settings.builder().put("location", "backup").build());
@@ -46,16 +55,25 @@ public void testRunningSnapshotCompleteAfterUpgrade() throws Exception {
             }
             flush(indexName, true);
             // Signal shutdown to prevent snapshot from being completed
-            putShutdownMetadata(nodeIds);
+            putShutdownMetadata(nodeIdToNodeNames.keySet());
             createSnapshot(repositoryName, snapshotName, false);
             assertRunningSnapshot(repositoryName, snapshotName);
         } else {
             if (isUpgradedCluster()) {
-                deleteShutdownMetadata(nodeIds);
-                assertNoShutdownMetadata(nodeIds);
+                deleteShutdownMetadata(nodeIdToNodeNames.keySet());
+                assertNoShutdownMetadata(nodeIdToNodeNames.keySet());
                 ensureGreen(indexName);
                 assertBusy(() -> assertCompletedSnapshot(repositoryName, snapshotName));
             } else {
+                if (isFirstMixedCluster()) {
+                    final var upgradedNodeIds = nodeIdToNodeNames.entrySet()
+                        .stream()
+                        .filter(entry -> "test-cluster-0".equals(entry.getValue()))
+                        .map(Map.Entry::getKey)
+                        .collect(Collectors.toUnmodifiableSet());
+                    assertThat(upgradedNodeIds, hasSize(1));
+                    deleteShutdownMetadata(upgradedNodeIds);
+                }
                 assertRunningSnapshot(repositoryName, snapshotName);
             }
         }
@@ -76,6 +94,7 @@ private void putShutdownMetadata(Collection<String> nodeIds) throws IOException
     private void deleteShutdownMetadata(Collection<String> nodeIds) throws IOException {
         for (String nodeId : nodeIds) {
             final Request request = new Request("DELETE", "/_nodes/" + nodeId + "/shutdown");
+            request.addParameter(IGNORE_RESPONSE_CODES_PARAM, "404");
             client().performRequest(request);
         }
     }