Fix testRetentionLeasesEstablishedWhenRelocatingPrimary (#52445)

Replace the current assertion with a more robust assertion.

Closes #52364

qa/full-cluster-restart/src/test/java/org/elasticsearch/upgrades/FullClusterRestartIT.java

@@ -34,7 +34,6 @@ import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.json.JsonXContent;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.index.IndexSettings;
-import org.elasticsearch.index.seqno.RetentionLeaseUtils;
 import org.elasticsearch.test.NotEqualMessageBuilder;
 import org.elasticsearch.test.rest.ESRestTestCase;
 import org.elasticsearch.test.rest.yaml.ObjectPath;
@@ -1237,7 +1236,7 @@ public class FullClusterRestartIT extends AbstractFullClusterRestartTestCase {
         assertFalse((Boolean) healthRsp.get("timed_out"));
     }
 
-    public void testPeerRecoveryRetentionLeases() throws IOException {
+    public void testPeerRecoveryRetentionLeases() throws Exception {
         if (isRunningAgainstOldCluster()) {
             XContentBuilder settings = jsonBuilder();
             settings.startObject();
@@ -1252,11 +1251,9 @@ public class FullClusterRestartIT extends AbstractFullClusterRestartTestCase {
             Request createIndex = new Request("PUT", "/" + index);
             createIndex.setJsonEntity(Strings.toString(settings));
             client().performRequest(createIndex);
-            ensureGreen(index);
-        } else {
-            ensureGreen(index);
-            RetentionLeaseUtils.assertAllCopiesHavePeerRecoveryRetentionLeases(client(), index);
         }
+        ensureGreen(index);
+        ensurePeerRecoveryRetentionLeasesRenewedAndSynced(index);
     }
 
     /**

qa/rolling-upgrade/src/test/java/org/elasticsearch/upgrades/RecoveryIT.java

@@ -33,7 +33,6 @@ import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.AbstractRunnable;
 import org.elasticsearch.common.xcontent.support.XContentMapValues;
 import org.elasticsearch.index.IndexSettings;
-import org.elasticsearch.index.seqno.RetentionLeaseUtils;
 import org.elasticsearch.test.rest.yaml.ObjectPath;
 import org.hamcrest.Matchers;
 
@@ -339,9 +338,7 @@ public class RecoveryIT extends AbstractRollingTestCase {
             }
         }
         ensureGreen(index);
-        if (CLUSTER_TYPE == ClusterType.UPGRADED) {
-            assertBusy(() -> RetentionLeaseUtils.assertAllCopiesHavePeerRecoveryRetentionLeases(client(), index));
-        }
+        ensurePeerRecoveryRetentionLeasesRenewedAndSynced(index);
     }
 
     public void testRetentionLeasesEstablishedWhenRelocatingPrimary() throws Exception {
@@ -384,16 +381,14 @@ public class RecoveryIT extends AbstractRollingTestCase {
                     final Request putSettingsRequest = new Request("PUT", "/" + index + "/_settings");
                     putSettingsRequest.setJsonEntity("{\"index.routing.allocation.exclude._name\":\"" + oldNodeName + "\"}");
                     assertOK(client().performRequest(putSettingsRequest));
-                    ensureGreen(index);
-                    assertBusy(() -> RetentionLeaseUtils.assertAllCopiesHavePeerRecoveryRetentionLeases(client(), index));
-                } else {
-                    ensureGreen(index);
                 }
+                ensureGreen(index);
+                ensurePeerRecoveryRetentionLeasesRenewedAndSynced(index);
                 break;
 
             case UPGRADED:
                 ensureGreen(index);
-                assertBusy(() -> RetentionLeaseUtils.assertAllCopiesHavePeerRecoveryRetentionLeases(client(), index));
+                ensurePeerRecoveryRetentionLeasesRenewedAndSynced(index);
                 break;
         }
     }

test/framework/src/main/java/org/elasticsearch/index/seqno/RetentionLeaseUtils.java

@@ -18,27 +18,11 @@
  */
 package org.elasticsearch.index.seqno;
 
-import org.elasticsearch.client.Request;
-import org.elasticsearch.client.RestClient;
-import org.elasticsearch.cluster.routing.RecoverySource;
-import org.elasticsearch.cluster.routing.ShardRouting;
-import org.elasticsearch.cluster.routing.UnassignedInfo;
-import org.elasticsearch.index.shard.ShardId;
-import org.elasticsearch.test.rest.yaml.ObjectPath;
-import org.junit.Assert;
-
-import java.io.IOException;
-import java.util.HashSet;
 import java.util.LinkedHashMap;
-import java.util.List;
 import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-import static org.hamcrest.Matchers.hasItems;
-
 public class RetentionLeaseUtils {
 
     private RetentionLeaseUtils() {
@@ -61,39 +45,4 @@ public class RetentionLeaseUtils {
                 },
                 LinkedHashMap::new));
     }
-
-    /**
-     * Asserts that every copy of every shard of the given index has a peer recovery retention lease according to the stats exposed by the
-     * REST API
-     */
-    public static void assertAllCopiesHavePeerRecoveryRetentionLeases(RestClient restClient, String index) throws IOException {
-        final Request statsRequest = new Request("GET", "/" + index + "/_stats");
-        statsRequest.addParameter("level", "shards");
-        final Map<?, ?> shardsStats = ObjectPath.createFromResponse(restClient.performRequest(statsRequest))
-            .evaluate("indices." + index + ".shards");
-        for (Map.Entry<?, ?> shardCopiesEntry : shardsStats.entrySet()) {
-            final List<?> shardCopiesList = (List<?>) shardCopiesEntry.getValue();
-
-            final Set<String> expectedLeaseIds = new HashSet<>();
-            for (Object shardCopyStats : shardCopiesList) {
-                final String nodeId
-                    = Objects.requireNonNull((String) ((Map<?, ?>) (((Map<?, ?>) shardCopyStats).get("routing"))).get("node"));
-                expectedLeaseIds.add(ReplicationTracker.getPeerRecoveryRetentionLeaseId(
-                    ShardRouting.newUnassigned(new ShardId("_na_", "test", 0), false, RecoverySource.PeerRecoverySource.INSTANCE,
-                        new UnassignedInfo(UnassignedInfo.Reason.INDEX_CREATED, "test")).initialize(nodeId, null, 0L)));
-            }
-
-            final Set<String> actualLeaseIds = new HashSet<>();
-            for (Object shardCopyStats : shardCopiesList) {
-                final List<?> leases
-                    = (List<?>) ((Map<?, ?>) (((Map<?, ?>) shardCopyStats).get("retention_leases"))).get("leases");
-                for (Object lease : leases) {
-                    actualLeaseIds.add(Objects.requireNonNull((String) (((Map<?, ?>) lease).get("id"))));
-                }
-            }
-            Assert.assertThat("[" + index + "][" + shardCopiesEntry.getKey() + "] has leases " + actualLeaseIds
-                    + " but expected " + expectedLeaseIds,
-                actualLeaseIds, hasItems(expectedLeaseIds.toArray(new String[0])));
-        }
-    }
 }