content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.kafka.controller;
  
  import org.apache.kafka.common.DirectoryId;
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.common.metadata.PartitionChangeRecord;
  import org.apache.kafka.common.protocol.types.TaggedFields;
  import org.apache.kafka.common.requests.AlterPartitionRequest;
  import org.apache.kafka.controller.PartitionChangeBuilder.ElectionResult;
  import org.apache.kafka.metadata.LeaderRecoveryState;
  import org.apache.kafka.metadata.PartitionRegistration;
  import org.apache.kafka.metadata.Replicas;
  import org.apache.kafka.metadata.placement.DefaultDirProvider;
  import org.apache.kafka.server.common.ApiMessageAndVersion;
  import org.apache.kafka.server.common.MetadataVersion;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.Timeout;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.Arguments;
  import org.junit.jupiter.params.provider.MethodSource;
  import org.junit.jupiter.params.provider.ValueSource;
  
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.function.IntPredicate;
  import java.util.stream.IntStream;
  import java.util.stream.Stream;
  
  import static org.apache.kafka.controller.PartitionChangeBuilder.Election;
  import static org.apache.kafka.controller.PartitionChangeBuilder.changeRecordIsNoOp;
  import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER;
  import static org.apache.kafka.metadata.LeaderConstants.NO_LEADER_CHANGE;
  import static org.apache.kafka.metadata.placement.PartitionAssignmentTest.partitionAssignment;
  import static org.junit.jupiter.api.Assertions.assertArrayEquals;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  
  
  @Timeout(value = 40)
  public class PartitionChangeBuilderTest {
      private static final DefaultDirProvider DEFAULT_DIR_PROVIDER = brokerId -> DirectoryId.UNASSIGNED;
  
      private static Stream<Arguments> partitionChangeRecordVersions() {
          return IntStream.range(PartitionChangeRecord.LOWEST_SUPPORTED_VERSION, PartitionChangeRecord.HIGHEST_SUPPORTED_VERSION + 1).mapToObj(version -> Arguments.of((short) version));
      }
  
      @Test
      public void testChangeRecordIsNoOp() {
          /* If the next few checks fail please update them based on the latest schema and make sure
           * to update changeRecordIsNoOp to take into account the new schema or tagged fields.
           */
          // Check that the supported versions haven't changed
          assertEquals(2, PartitionChangeRecord.HIGHEST_SUPPORTED_VERSION);
          assertEquals(0, PartitionChangeRecord.LOWEST_SUPPORTED_VERSION);
          // For the latest version check that the number of tagged fields hasn't changed
          TaggedFields taggedFields = (TaggedFields) PartitionChangeRecord.SCHEMA_0.get(2).def.type;
          assertEquals(6, taggedFields.numFields());
  
          assertTrue(changeRecordIsNoOp(new PartitionChangeRecord()));
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().setLeader(1)));
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().
              setIsr(List.of(1, 2, 3))));
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().
              setRemovingReplicas(List.of(1))));
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().
              setAddingReplicas(List.of(4))));
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().
                  setEligibleLeaderReplicas(List.of(5))));
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().
                  setLastKnownElr(List.of(6))));
          assertFalse(
              changeRecordIsNoOp(
                  new PartitionChangeRecord()
                    .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED.value())
              )
          );
          assertFalse(changeRecordIsNoOp(new PartitionChangeRecord().setDirectories(List.of(
                  Uuid.fromString("5JwD0VNXRV2Wr9CCON38Tw"),
                  Uuid.fromString("zpL1bRzTQXmmgdxlLHOWuw"),
                  Uuid.fromString("6iGUpAkHQXC6bY0FTcPRDw")
          ))));
      }
  
      private static final PartitionRegistration FOO = new PartitionRegistration.Builder().
          setReplicas(new int[] {2, 1, 3}).
          setDirectories(new Uuid[]{
              Uuid.fromString("dpdvA5AZSWySmnPFTnu5Kw"),
              Uuid.fromString("V60B3cglScq3Xk8BX1NxAQ"),
              DirectoryId.UNASSIGNED,
          }).
          setIsr(new int[] {2, 1, 3}).
          setLeader(1).
          setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
          setLeaderEpoch(100).
          setPartitionEpoch(200).
          build();
  
      private static final Uuid FOO_ID = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
  
      private static MetadataVersion metadataVersionForPartitionChangeRecordVersion(short version) {
          switch (version) {
              case (short) 0:
                  return MetadataVersion.IBP_3_7_IV0;
              case (short) 1:
                  return MetadataVersion.IBP_3_7_IV2;
              case (short) 2:
                  return MetadataVersion.IBP_4_0_IV1;
              default:
                  throw new RuntimeException("Unknown PartitionChangeRecord version " + version);
          }
      }
  
      private static PartitionChangeBuilder createFooBuilder(MetadataVersion metadataVersion) {
          return new PartitionChangeBuilder(FOO,
                  FOO_ID,
                  0,
                  r -> r != 3,
                  metadataVersion,
                  2).
                  setEligibleLeaderReplicasEnabled(metadataVersion.isElrSupported()).
                  setDefaultDirProvider(DEFAULT_DIR_PROVIDER);
      }
  
      private static PartitionChangeBuilder createFooBuilder(short partitionChangeRecordVersion) {
          return createFooBuilder(metadataVersionForPartitionChangeRecordVersion(partitionChangeRecordVersion));
      }
  
      private static final PartitionRegistration BAR = new PartitionRegistration.Builder().
          setReplicas(new int[] {1, 2, 3, 4}).
          setDirectories(new Uuid[] {
              DirectoryId.UNASSIGNED,
              Uuid.fromString("X5FnAcIgTheWgTMzeO5WHw"),
              Uuid.fromString("GtrcdoSOTm2vFMGFeZq0eg"),
              Uuid.fromString("YcOqPw5ARmeKr1y9W3AkFw"),
          }).
          setIsr(new int[] {1, 2, 3}).
          setRemovingReplicas(new int[] {1}).
          setAddingReplicas(new int[] {4}).
          setLeader(1).
          setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
          setLeaderEpoch(100).
          setPartitionEpoch(200).
          build();
  
      private static final Uuid BAR_ID = Uuid.fromString("LKfUsCBnQKekvL9O5dY9nw");
  
      private static boolean isElrEnabled(short partitionChangeRecordVersion) {
          return partitionChangeRecordVersion >= 2;
      }
  
      private static PartitionChangeBuilder createBarBuilder(short version) {
          return new PartitionChangeBuilder(BAR,
                  BAR_ID,
                  0,
                  r -> r != 3,
                  metadataVersionForPartitionChangeRecordVersion(version),
                  2).
                  setEligibleLeaderReplicasEnabled(isElrEnabled(version)).
                  setDefaultDirProvider(DEFAULT_DIR_PROVIDER);
      }
  
      private static final PartitionRegistration BAZ = new PartitionRegistration.Builder().
          setReplicas(new int[] {2, 1, 3}).
          setDirectories(new Uuid[] {
              Uuid.fromString("ywnfFpTBTbOsFdZ6uAdOmw"),
              Uuid.fromString("Th0x70ecRbWvZNNV33jyRA"),
              Uuid.fromString("j216tuSoQsC9JFd1Z5ZP6w"),
          }).
          setIsr(new int[] {1, 3}).
          setLeader(3).
          setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
          setLeaderEpoch(100).
          setPartitionEpoch(200).
          build();
  
      private static final Uuid BAZ_ID = Uuid.fromString("wQzt5gkSTwuQNXZF5gIw7A");
  
      private static PartitionChangeBuilder createBazBuilder(short version) {
          return new PartitionChangeBuilder(BAZ,
                  BAZ_ID,
                  0,
                  __ -> true,
                  metadataVersionForPartitionChangeRecordVersion(version),
                  2).
                  setEligibleLeaderReplicasEnabled(isElrEnabled(version)).
                  setDefaultDirProvider(DEFAULT_DIR_PROVIDER);
      }
  
      private static final PartitionRegistration OFFLINE_WITHOUT_ELR = new PartitionRegistration.Builder().
          setReplicas(new int[] {2, 1, 3}).
          setDirectories(new Uuid[]{
             Uuid.fromString("iYGgiDV5Sb2EtH6hbgYnCA"),
             Uuid.fromString("XI2t4qAUSkGlLZSKeEVf8g"),
             Uuid.fromString("eqRW24kIRlitzQFzmovE0Q")
          }).
          setIsr(new int[] {3}).
          setLeader(-1).
          setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
          setLeaderEpoch(100).
          setPartitionEpoch(200).
          build();
  
      private static final PartitionRegistration OFFLINE_WITH_ELR = new PartitionRegistration.Builder().
              setReplicas(new int[] {2, 1, 3}).
              setDirectories(new Uuid[]{
                      Uuid.fromString("CQEqt7trRrmqyNxUT1CY0g"),
                      Uuid.fromString("59Mb9smoSsC0bGUP2FYV8A"),
                      Uuid.fromString("LBTmsCVJREqJuIEtwqxRDg")
              }).
              setElr(new int[] {3}).
              setIsr(new int[] {}).
              setLastKnownElr(new int[] {2}).
              setLeader(-1).
              setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
              setLeaderEpoch(100).
              setPartitionEpoch(200).
              build();
  
      private static final Uuid OFFLINE_ID = Uuid.fromString("LKfUsCBnQKekvL9O5dY9nw");
  
      private static PartitionChangeBuilder createOfflineBuilder(short partitionChangeRecordVersion) {
          MetadataVersion metadataVersion =
              metadataVersionForPartitionChangeRecordVersion(partitionChangeRecordVersion);
          if (metadataVersion.isElrSupported()) {
              return new PartitionChangeBuilder(OFFLINE_WITH_ELR, OFFLINE_ID, 0, r -> r == 1,
                      metadataVersion, 2).
                       setEligibleLeaderReplicasEnabled(true).
                       setDefaultDirProvider(DEFAULT_DIR_PROVIDER);
          } else {
              return new PartitionChangeBuilder(OFFLINE_WITHOUT_ELR, OFFLINE_ID, 0, r -> r == 1,
                      metadataVersion, 2).
                       setEligibleLeaderReplicasEnabled(false).
                       setDefaultDirProvider(DEFAULT_DIR_PROVIDER);
          }
      }
  
      private static void assertElectLeaderEquals(PartitionChangeBuilder builder,
                                                 int expectedNode,
                                                 boolean expectedUnclean) {
          ElectionResult electionResult = builder.electLeader();
          assertEquals(expectedNode, electionResult.node);
          assertEquals(expectedUnclean, electionResult.unclean);
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testElectLeader(short version) {
          assertElectLeaderEquals(createFooBuilder(version).setElection(Election.PREFERRED), 2, false);
          assertElectLeaderEquals(createFooBuilder(version), 1, false);
          assertElectLeaderEquals(createFooBuilder(version).setElection(Election.UNCLEAN), 1, false);
          assertElectLeaderEquals(createFooBuilder(version)
              .setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(1, 3))), 1, false);
          assertElectLeaderEquals(createFooBuilder(version).setElection(Election.UNCLEAN)
              .setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(1, 3))), 1, false);
          assertElectLeaderEquals(createFooBuilder(version)
              .setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(3))), NO_LEADER, false);
          assertElectLeaderEquals(createFooBuilder(version).setElection(Election.UNCLEAN).
              setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(3))), 2, true);
          assertElectLeaderEquals(
              createFooBuilder(version).setElection(Election.UNCLEAN)
                  .setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(4))).setTargetReplicas(List.of(2, 1, 3, 4)),
              4,
              false
          );
  
          assertElectLeaderEquals(createBazBuilder(version).setElection(Election.PREFERRED), 3, false);
          assertElectLeaderEquals(createBazBuilder(version), 3, false);
          assertElectLeaderEquals(createBazBuilder(version).setElection(Election.UNCLEAN), 3, false);
      }
  
      private static void testTriggerLeaderEpochBumpIfNeeded(
          PartitionChangeBuilder builder,
          PartitionChangeRecord record,
          int expectedLeader
      ) {
          builder.triggerLeaderEpochBumpForReplicaReassignmentIfNeeded(record);
          record.setIsr(builder.targetIsr());
          builder.triggerLeaderEpochBumpForIsrShrinkIfNeeded(record);
          assertEquals(expectedLeader, record.leader());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testNoLeaderEpochBumpIfNothingChanged(short version) {
          testTriggerLeaderEpochBumpIfNeeded(createFooBuilder(version),
              new PartitionChangeRecord(),
              NO_LEADER_CHANGE);
      }
  
      /**
       * Test that shrinking the ISR doesn't increase the leader epoch in later MVs.
       */
      @ParameterizedTest
      @ValueSource(strings = {"3.6-IV0", "3.7-IV2", "4.0-IV0"})
      public void testNoLeaderEpochBumpOnIsrShrink(String metadataVersionString) {
          MetadataVersion metadataVersion = MetadataVersion.fromVersionString(metadataVersionString);
          testTriggerLeaderEpochBumpIfNeeded(
              createFooBuilder(metadataVersion).setTargetIsrWithBrokerStates(
                  AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(2, 1))),
              new PartitionChangeRecord(),
              NO_LEADER_CHANGE);
      }
  
      /**
       * Test that shrinking the ISR does increase the leader epoch in earlier MVs.
       * See KAFKA-15021 for details.
       */
      @ParameterizedTest
      @ValueSource(strings = {"3.4-IV0", "3.5-IV2"})
      public void testLeaderEpochBumpOnIsrShrink(String metadataVersionString) {
          MetadataVersion metadataVersion = MetadataVersion.fromVersionString(metadataVersionString);
          testTriggerLeaderEpochBumpIfNeeded(
              createFooBuilder(metadataVersion).setTargetIsrWithBrokerStates(
                  AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(2, 1))),
              new PartitionChangeRecord(),
              1);
      }
  
      /**
       * Test that expanding the ISR doesn't increase the leader epoch.
       */
      @ParameterizedTest
      @ValueSource(strings = {"3.4-IV0", "3.5-IV2", "3.6-IV0", "3.7-IV2", "4.0-IV0"})
      public void testNoLeaderEpochBumpOnIsrExpansion(String metadataVersionString) {
          MetadataVersion metadataVersion = MetadataVersion.fromVersionString(metadataVersionString);
          testTriggerLeaderEpochBumpIfNeeded(
              createFooBuilder(metadataVersion).setTargetIsrWithBrokerStates(
                  AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(2, 1, 3, 4))),
              new PartitionChangeRecord(),
              NO_LEADER_CHANGE);
      }
  
      /**
       * Test that changing the replica set such that not all the old replicas remain
       * always results in a leader epoch increase.
       */
      @ParameterizedTest
      @ValueSource(strings = {"3.4-IV0", "3.5-IV2", "3.6-IV0", "3.7-IV2", "4.0-IV0"})
      public void testLeaderEpochBumpOnNewReplicaSetDisjoint(String metadataVersionString) {
          MetadataVersion metadataVersion = MetadataVersion.fromVersionString(metadataVersionString);
          testTriggerLeaderEpochBumpIfNeeded(
              createFooBuilder(metadataVersion).setTargetReplicas(List.of(2, 1, 4)),
              new PartitionChangeRecord(),
              1);
      }
  
      /**
       * Regression test for KAFKA-16624. Tests that when targetIsr is the empty list, but we
       * cannot actually change the ISR, triggerLeaderEpochBumpForIsrShrinkIfNeeded does not engage.
       */
      @ParameterizedTest
      @ValueSource(strings = {"3.4-IV0", "3.5-IV2", "3.6-IV0", "3.7-IV2"})
      public void testNoLeaderEpochBumpOnEmptyTargetIsr(String metadataVersionString) {
          MetadataVersion metadataVersion = MetadataVersion.fromVersionString(metadataVersionString);
          PartitionRegistration partition = new PartitionRegistration.Builder().
              setReplicas(new int[] {2}).
              setDirectories(new Uuid[]{
                  Uuid.fromString("dpdvA5AZSWySmnPFTnu5Kw")
              }).
              setIsr(new int[] {2}).
              setLeader(2).
              setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
              setLeaderEpoch(100).
              setPartitionEpoch(200).
              build();
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition,
              FOO_ID,
              0,
              r -> true,
              metadataVersion,
              2).
              setEligibleLeaderReplicasEnabled(metadataVersion.isElrSupported()).
              setDefaultDirProvider(DEFAULT_DIR_PROVIDER).
              setTargetReplicas(List.of());
          PartitionChangeRecord record = new PartitionChangeRecord();
          builder.triggerLeaderEpochBumpForIsrShrinkIfNeeded(record);
          assertEquals(NO_LEADER_CHANGE, record.leader());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testNoChange(short version) {
          assertEquals(Optional.empty(), createFooBuilder(version).build());
          assertEquals(Optional.empty(), createFooBuilder(version).setElection(Election.UNCLEAN).build());
          assertEquals(Optional.empty(), createBarBuilder(version).build());
          assertEquals(Optional.empty(), createBarBuilder(version).setElection(Election.UNCLEAN).build());
          assertEquals(Optional.empty(), createBazBuilder(version).setElection(Election.PREFERRED).build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testIsrChangeDoesntBumpLeaderEpoch(short version) {
          // Changing the ISR should not cause the leader epoch to increase
          assertEquals(
              // Expected
              Optional.of(
                  new ApiMessageAndVersion(
                      new PartitionChangeRecord()
                        .setTopicId(FOO_ID)
                        .setPartitionId(0)
                        .setIsr(List.of(2, 1)),
                      version
                  )
              ),
              // Actual
              createFooBuilder(version)
                .setTargetIsrWithBrokerStates(
                    AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(2, 1))
                )
                .build()
          );
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testIsrChangeAndLeaderChange(short version) {
          assertEquals(Optional.of(new ApiMessageAndVersion(new PartitionChangeRecord().
                  setTopicId(FOO_ID).
                  setPartitionId(0).
                  setIsr(List.of(2, 3)).
                  setLeader(2), version)),
              createFooBuilder(version).setTargetIsrWithBrokerStates(AlterPartitionRequest.
                  newIsrToSimpleNewIsrWithBrokerEpochs(List.of(2, 3))).build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testReassignmentRearrangesReplicas(short version) {
          PartitionChangeRecord expectedRecord = new PartitionChangeRecord().
                  setTopicId(FOO_ID).
                  setPartitionId(0).
                  setReplicas(List.of(3, 2, 1));
          if (version >= 1) {
              Map<Integer, Uuid> dirs = DirectoryId.createAssignmentMap(FOO.replicas, FOO.directories);
              expectedRecord.setDirectories(List.of(dirs.get(3), dirs.get(2), dirs.get(1)));
          }
          assertEquals(Optional.of(new ApiMessageAndVersion(expectedRecord, version)),
              createFooBuilder(version).setTargetReplicas(List.of(3, 2, 1)).build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testIsrEnlargementCompletesReassignment(short version) {
          PartitionChangeRecord expectedRecord = new PartitionChangeRecord().
                  setTopicId(BAR_ID).
                  setPartitionId(0).
                  setReplicas(List.of(2, 3, 4)).
                  setIsr(List.of(2, 3, 4)).
                  setLeader(2).
                  setRemovingReplicas(List.of()).
                  setAddingReplicas(List.of());
          if (version >= 1) {
              Map<Integer, Uuid> dirs = DirectoryId.createAssignmentMap(BAR.replicas, BAR.directories);
              expectedRecord.setDirectories(List.of(dirs.get(2), dirs.get(3), dirs.get(4)));
          }
          assertEquals(Optional.of(new ApiMessageAndVersion(expectedRecord, version)),
              createBarBuilder(version).setTargetIsrWithBrokerStates(AlterPartitionRequest.
                  newIsrToSimpleNewIsrWithBrokerEpochs(List.of(1, 2, 3, 4))).build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testRevertReassignment(short version) {
          PartitionReassignmentRevert revert = new PartitionReassignmentRevert(BAR);
          assertEquals(List.of(1, 2, 3), revert.replicas());
          assertEquals(List.of(1, 2, 3), revert.isr());
          PartitionChangeRecord expectedRecord = new PartitionChangeRecord().
                  setTopicId(BAR_ID).
                  setPartitionId(0).
                  setReplicas(List.of(1, 2, 3)).
                  setLeader(1).
                  setRemovingReplicas(List.of()).
                  setAddingReplicas(List.of());
          if (version >= 1) {
              Map<Integer, Uuid> dirs = DirectoryId.createAssignmentMap(BAR.replicas, BAR.directories);
              expectedRecord.setDirectories(List.of(dirs.get(1), dirs.get(2), dirs.get(3)));
          }
          assertEquals(Optional.of(new ApiMessageAndVersion(expectedRecord, version)),
              createBarBuilder(version).
                  setTargetReplicas(revert.replicas()).
                  setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(revert.isr())).
                  setTargetRemoving(List.of()).
                  setTargetAdding(List.of()).
                  build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testRemovingReplicaReassignment(short version) {
          PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
              partitionAssignment(Replicas.toList(FOO.replicas)), partitionAssignment(List.of(1, 2)));
          assertEquals(List.of(3), replicas.removing());
          assertEquals(List.of(), replicas.adding());
          assertEquals(List.of(1, 2, 3), replicas.replicas());
          PartitionChangeRecord expectedRecord = new PartitionChangeRecord().
                  setTopicId(FOO_ID).
                  setPartitionId(0).
                  setReplicas(List.of(1, 2)).
                  setIsr(List.of(2, 1)).
                  setLeader(1);
          if (version >= 1) {
              Map<Integer, Uuid> dirs = DirectoryId.createAssignmentMap(FOO.replicas, FOO.directories);
              expectedRecord.setDirectories(List.of(dirs.get(1), dirs.get(2)));
          }
          assertEquals(Optional.of(new ApiMessageAndVersion(expectedRecord, version)),
              createFooBuilder(version).
                  setTargetReplicas(replicas.replicas()).
                  setTargetRemoving(replicas.removing()).
                  build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testAddingReplicaReassignment(short version) {
          PartitionReassignmentReplicas replicas = new PartitionReassignmentReplicas(
              partitionAssignment(Replicas.toList(FOO.replicas)), partitionAssignment(List.of(1, 2, 3, 4)));
          assertEquals(List.of(), replicas.removing());
          assertEquals(List.of(4), replicas.adding());
          assertEquals(List.of(1, 2, 3, 4), replicas.replicas());
          PartitionChangeRecord expectedRecord = new PartitionChangeRecord().
                  setTopicId(FOO_ID).
                  setPartitionId(0).
                  setReplicas(List.of(1, 2, 3, 4)).
                  setAddingReplicas(List.of(4));
          if (version >= 1) {
              Map<Integer, Uuid> dirs = DirectoryId.createAssignmentMap(FOO.replicas, FOO.directories);
              expectedRecord.setDirectories(List.of(dirs.get(1), dirs.get(2), dirs.get(3), DirectoryId.UNASSIGNED));
          }
          assertEquals(Optional.of(new ApiMessageAndVersion(expectedRecord, version)),
              createFooBuilder(version).
                  setTargetReplicas(replicas.replicas()).
                  setTargetAdding(replicas.adding()).
                  build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testUncleanLeaderElection(short version) {
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(
              new PartitionChangeRecord()
                  .setTopicId(FOO_ID)
                  .setPartitionId(0)
                  .setIsr(List.of(2))
                  .setLeader(2)
                  .setLeaderRecoveryState(LeaderRecoveryState.RECOVERING.value()),
              version
          );
          assertEquals(
              Optional.of(expectedRecord),
              createFooBuilder(version).setElection(Election.UNCLEAN)
                  .setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(3))).build()
          );
  
          PartitionChangeRecord record = new PartitionChangeRecord()
              .setTopicId(OFFLINE_ID)
              .setPartitionId(0)
              .setIsr(List.of(1))
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERING.value());
  
          if (version >= 2) {
              // The test partition has ELR, so unclean election will clear these fields.
              record.setEligibleLeaderReplicas(List.of())
                  .setLastKnownElr(List.of());
          }
  
          expectedRecord = new ApiMessageAndVersion(record, version);
          assertEquals(
              Optional.of(expectedRecord),
              createOfflineBuilder(version).setElection(Election.UNCLEAN).build()
          );
  
          assertEquals(
              Optional.of(expectedRecord),
              createOfflineBuilder(version).setElection(Election.UNCLEAN)
                  .setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(2))).build()
          );
      }
  
      @Test
      public void testChangeInLeadershipDoesNotChangeRecoveryState() {
          final byte noChange = (byte) -1;
          int leaderId = 1;
          LeaderRecoveryState recoveryState = LeaderRecoveryState.RECOVERING;
          PartitionRegistration registration = new PartitionRegistration.Builder().
              setReplicas(new int[] {leaderId, leaderId + 1, leaderId + 2}).
              setDirectories(new Uuid[] {
                      Uuid.fromString("1sF6XXLkSN2LtDums7CJ8Q"),
                      Uuid.fromString("iaBBVsoHQR6NDKXwliKMqw"),
                      Uuid.fromString("sHaBwjdrR2S3bL4E1RKC8Q")
              }).
              setIsr(new int[] {leaderId}).
              setLeader(leaderId).
              setLeaderRecoveryState(recoveryState).
              setLeaderEpoch(100).
              setPartitionEpoch(200).
              build();
  
          MetadataVersion metadataVersion = MetadataVersion.MINIMUM_VERSION;
  
          // Change the partition so that there is no leader
          PartitionChangeBuilder offlineBuilder = new PartitionChangeBuilder(
              registration,
              FOO_ID,
              0,
              brokerId -> false,
              metadataVersion,
              2
          );
          // Set the target ISR to empty to indicate that the last leader is offline
          offlineBuilder.setTargetIsrWithBrokerStates(List.of());
  
          // The partition should stay as recovering
          PartitionChangeRecord changeRecord = (PartitionChangeRecord) offlineBuilder
              .build()
              .get()
              .message();
          assertEquals(noChange, changeRecord.leaderRecoveryState());
          assertEquals(NO_LEADER, changeRecord.leader());
  
          registration = registration.merge(changeRecord);
  
          assertEquals(NO_LEADER, registration.leader);
          assertEquals(leaderId, registration.isr[0]);
          assertEquals(recoveryState, registration.leaderRecoveryState);
  
          // Bring the leader back online
          PartitionChangeBuilder onlineBuilder = new PartitionChangeBuilder(
              registration,
              FOO_ID,
              0,
              brokerId -> true,
              metadataVersion,
              2
          );
  
          // The only broker in the ISR is elected leader and stays in the recovering
          changeRecord = (PartitionChangeRecord) onlineBuilder.build().get().message();
          assertEquals(noChange, changeRecord.leaderRecoveryState());
  
          registration = registration.merge(changeRecord);
  
          assertEquals(leaderId, registration.leader);
          assertEquals(leaderId, registration.isr[0]);
          assertEquals(recoveryState, registration.leaderRecoveryState);
      }
  
      @Test
      void testUncleanSetsLeaderRecoveringState() {
          int leaderId = 1;
          PartitionRegistration registration = new PartitionRegistration.Builder().
              setReplicas(new int[] {leaderId, leaderId + 1, leaderId + 2}).
              setDirectories(new Uuid[] {
                  Uuid.fromString("uYpxts0pS4K4bk5XOoXB4g"),
                  Uuid.fromString("kS6fHEqwRYucduWkmvsevw"),
                  Uuid.fromString("De9RqRThQRGjKg3i3yzUxA")
              }).
              setIsr(new int[] {leaderId + 1, leaderId + 2}).
              setLeader(NO_LEADER).
              setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
              setLeaderEpoch(100).
              setPartitionEpoch(200).
              build();
  
          MetadataVersion metadataVersion = MetadataVersion.MINIMUM_VERSION;
  
          // Change the partition using unclean leader election
          PartitionChangeBuilder onlineBuilder = new PartitionChangeBuilder(
              registration,
              FOO_ID,
              0,
              brokerId -> brokerId == leaderId,
              metadataVersion,
              2
          ).setElection(Election.UNCLEAN);
          // The partition should stay as recovering
          PartitionChangeRecord changeRecord = (PartitionChangeRecord) onlineBuilder
              .build()
              .get()
              .message();
  
          assertEquals(LeaderRecoveryState.RECOVERING.value(), changeRecord.leaderRecoveryState());
          assertEquals(leaderId, changeRecord.leader());
          assertEquals(1, changeRecord.isr().size());
          assertEquals(leaderId, changeRecord.isr().get(0));
  
          registration = registration.merge(changeRecord);
  
          assertEquals(leaderId, registration.leader);
          assertEquals(leaderId, registration.isr[0]);
          assertEquals(LeaderRecoveryState.RECOVERING, registration.leaderRecoveryState);
      }
  
      @Test
      public void testStoppedLeaderIsDemotedAfterReassignmentCompletesEvenIfNoNewEligibleLeaders() {
          // Set up PartitionRegistration as if there's an ongoing reassignment from [0, 1] to [2, 3]
          int[] replicas = new int[] {2, 3, 0, 1};
          Uuid[] directories = {
                  Uuid.fromString("XCBQClkBSZyphD87QUXzDA"),
                  Uuid.fromString("Or2Rp9tTQOSVuy12hsfmTA"),
                  Uuid.fromString("pThsodMNSwGvljTfc1RNVQ"),
                  Uuid.fromString("d8CGoNJmS5mJdF20tc8P7g")
          };
          // The ISR starts off with the old replicas
          int[] isr = new int[] {0, 1};
          // We're removing [0, 1]
          int[] removingReplicas = new int[] {0, 1};
          // And adding [2, 3]
          int[] addingReplicas = new int[] {2, 3};
          // The leader is 0, one of the replicas we're removing
          int leader = 0;
          LeaderRecoveryState leaderRecoveryState = LeaderRecoveryState.RECOVERED;
          int leaderEpoch = 0;
          int partitionEpoch = 0;
          PartitionRegistration part = new PartitionRegistration.Builder().
              setReplicas(replicas).
              setDirectories(directories).
              setIsr(isr).
              setRemovingReplicas(removingReplicas).
              setAddingReplicas(addingReplicas).
              setLeader(leader).
              setLeaderRecoveryState(leaderRecoveryState).
              setLeaderEpoch(leaderEpoch).
              setPartitionEpoch(partitionEpoch).
              build();
  
          Uuid topicId = Uuid.randomUuid();
          // Always return false for valid leader. This is so none of the new replicas are valid leaders. This is so we
          // test what happens when the previous leader is a replica being "stopped" ie removed from the replicas list
          // and none of the adding replicas can be a leader. We want to make sure we do not leave the previous replica
          // being stopped as leader.
          IntPredicate isValidLeader = l -> false;
  
          PartitionChangeBuilder partitionChangeBuilder = new PartitionChangeBuilder(
              part,
              topicId,
              0,
              isValidLeader,
              MetadataVersion.MINIMUM_VERSION,
              2
          );
  
          // Before we build the new PartitionChangeBuilder, confirm the current leader is 0.
          assertEquals(0, part.leader);
          // The important part is that the new leader is NO_LEADER.
          assertEquals(Optional.of(new ApiMessageAndVersion(new PartitionChangeRecord().
                  setTopicId(topicId).
                  setPartitionId(0).
                  setReplicas(List.of(2, 3)).
                  setIsr(List.of(2, 3)).
                  setRemovingReplicas(List.of()).
                  setAddingReplicas(List.of()).
                  setLeader(NO_LEADER),
                  (short) 0)),
              partitionChangeBuilder.setTargetIsr(List.of(0, 1, 2, 3)).
                  build());
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testEligibleLeaderReplicas_IsrShrinkBelowMinISR(short version) {
          PartitionRegistration partition = new PartitionRegistration.Builder()
              .setReplicas(new int[] {1, 2, 3, 4})
              .setDirectories(new Uuid[] {
                  Uuid.fromString("NeQeLdHhSXi4tQGaFcszKA"),
                  Uuid.fromString("LsVrQZ73RSSuEWA8hhqQhg"),
                  Uuid.fromString("0IaY4zXKRR6sROgE8yHfnw"),
                  Uuid.fromString("1WxphfLCSZqMHKK4JMppuw")
              })
              .setIsr(new int[] {1, 2, 3, 4})
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
              .setLeaderEpoch(100)
              .setPartitionEpoch(200)
              .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> r != 3,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(isElrEnabled(version))
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setUseLastKnownLeaderInBalancedRecovery(false);
  
          // Update ISR to {1, 2}
          builder.setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(1, 2)));
  
          PartitionChangeRecord record = new PartitionChangeRecord()
              .setTopicId(topicId)
              .setPartitionId(0)
              .setIsr(List.of(1, 2))
              .setLeader(-2)
              .setLeaderRecoveryState(LeaderRecoveryState.NO_CHANGE);
          if (version >= 2) {
              record.setEligibleLeaderReplicas(List.of(3, 4));
          }
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(record, version);
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          if (version >= 2) {
              assertArrayEquals(new int[]{3, 4}, partition.elr, partition.toString());
              assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
          } else {
              assertEquals(0, partition.elr.length);
              assertEquals(0, partition.lastKnownElr.length);
          }
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testEligibleLeaderReplicas_IsrExpandAboveMinISR(short version) {
          PartitionRegistration partition = new PartitionRegistration.Builder()
              .setReplicas(new int[] {1, 2, 3, 4})
              .setDirectories(new Uuid[]{
                  Uuid.fromString("CWgRKBKkToGn1HKzNb2qqQ"),
                  Uuid.fromString("SCnk7zfSQSmlKqvV702d3A"),
                  Uuid.fromString("9tO0QHlJRhimjKfH8m9d8A"),
                  Uuid.fromString("JaaqVOxNT2OGVNCCIFA2JQ")
              })
              .setIsr(new int[] {1, 2})
              .setElr(new int[] {3})
              .setLastKnownElr(new int[] {4})
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
              .setLeaderEpoch(100)
              .setPartitionEpoch(200)
              .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
          // Min ISR is 3.
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> r != 3,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(isElrEnabled(version))
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setUseLastKnownLeaderInBalancedRecovery(false);
  
          builder.setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(1, 2, 3)));
          PartitionChangeRecord record = new PartitionChangeRecord()
              .setTopicId(topicId)
              .setPartitionId(0)
              .setIsr(List.of(1, 2, 3))
              .setLeader(-2)
              .setLeaderRecoveryState(LeaderRecoveryState.NO_CHANGE);
  
          // Both versions will set the elr and lastKnownElr as empty list.
          record.setEligibleLeaderReplicas(List.of())
              .setLastKnownElr(List.of());
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(record, version);
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          assertEquals(0, partition.elr.length);
          assertEquals(0, partition.lastKnownElr.length);
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testEligibleLeaderReplicas_IsrAddNewMemberNotInELR(short version) {
          PartitionRegistration partition = new PartitionRegistration.Builder()
              .setReplicas(new int[] {1, 2, 3, 4})
              .setDirectories(new Uuid[]{
                  Uuid.fromString("gPcIwlldQXikdUB3F4GB6w"),
                  Uuid.fromString("gFs7V8mKR66z8e5qwtjIMA"),
                  Uuid.fromString("zKHU2fwrRkuypqTgITl46g"),
                  Uuid.fromString("zEgmBBh8QJGqbBIvzvH7JA")
              })
              .setIsr(new int[] {1})
              .setElr(new int[] {3})
              .setLastKnownElr(new int[] {2})
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
              .setLeaderEpoch(100)
              .setPartitionEpoch(200)
              .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
          // Min ISR is 3.
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> r != 3,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(isElrEnabled(version))
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setUseLastKnownLeaderInBalancedRecovery(false);
  
          builder.setTargetIsrWithBrokerStates(AlterPartitionRequest.newIsrToSimpleNewIsrWithBrokerEpochs(List.of(1, 4)));
          PartitionChangeRecord record = new PartitionChangeRecord()
              .setTopicId(topicId)
              .setPartitionId(0)
              .setIsr(List.of(1, 4))
              .setLeader(-2)
              .setLeaderRecoveryState(LeaderRecoveryState.NO_CHANGE);
          if (version < 2) {
              record.setEligibleLeaderReplicas(List.of());
              record.setLastKnownElr(List.of());
          }
          // No change is expected to ELR/LastKnownElr.
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(record, version);
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          if (version >= 2) {
              assertArrayEquals(new int[]{3}, partition.elr, partition.toString());
              assertArrayEquals(new int[]{2}, partition.lastKnownElr, partition.toString());
          } else {
              assertArrayEquals(new int[]{}, partition.elr, partition.toString());
              assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
          }
      }
  
      @ParameterizedTest
      @MethodSource("partitionChangeRecordVersions")
      public void testEligibleLeaderReplicas_RemoveUncleanShutdownReplicasFromElr(short version) {
          PartitionRegistration partition = new PartitionRegistration.Builder()
              .setReplicas(new int[] {1, 2, 3, 4})
              .setDirectories(new Uuid[] {
                  Uuid.fromString("keB9ssIPRlibyVJT5FcBVA"),
                  Uuid.fromString("FhezfoReTSmHoKxi8wOIOg"),
                  Uuid.fromString("QHtFxu8LShm6RiyAP6PxYg"),
                  Uuid.fromString("tUJOMtvMQkGga30ydluvbQ")
              })
              .setIsr(new int[] {1})
              .setElr(new int[] {2, 3})
              .setLastKnownElr(new int[] {})
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
              .setLeaderEpoch(100)
              .setPartitionEpoch(200)
              .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
          // Min ISR is 3.
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> r != 3,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(isElrEnabled(version))
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setUseLastKnownLeaderInBalancedRecovery(false);
  
          builder.setUncleanShutdownReplicas(List.of(3));
  
          PartitionChangeRecord record = new PartitionChangeRecord()
              .setTopicId(topicId)
              .setPartitionId(0)
              .setLeader(-2)
              .setLeaderRecoveryState(LeaderRecoveryState.NO_CHANGE);
          if (version >= 2) {
              record.setEligibleLeaderReplicas(List.of(2))
                  .setLastKnownElr(List.of(3));
          } else {
              record.setEligibleLeaderReplicas(List.of());
          }
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(record, version);
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          if (version >= 2) {
              assertArrayEquals(new int[]{2}, partition.elr, partition.toString());
              assertArrayEquals(new int[]{3}, partition.lastKnownElr, partition.toString());
          } else {
              assertArrayEquals(new int[]{}, partition.elr, partition.toString());
              assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
          }
      }
  
      @Test
      public void testKeepsDirectoriesAfterReassignment() {
          PartitionRegistration registration = new PartitionRegistration.Builder().
                  setReplicas(new int[]{2, 1, 3}).
                  setDirectories(new Uuid[]{
                          Uuid.fromString("v1PVrX6uS5m8CByXlLfmWg"),
                          Uuid.fromString("iU2znv45Q9yQkOpkTSy3jA"),
                          Uuid.fromString("fM5NKyWTQHqEihjIkUl99Q")
                  }).
                  setIsr(new int[]{2, 1, 3}).
                  setLeader(1).
                  setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
                  setLeaderEpoch(100).
                  setPartitionEpoch(200).
                  build();
          Optional<ApiMessageAndVersion> built = new PartitionChangeBuilder(registration, FOO_ID,
                  0, r -> true, MetadataVersion.IBP_3_7_IV2, 2).
                  setTargetReplicas(List.of(3, 1, 5, 4)).
                  setDirectory(5, Uuid.fromString("RNJ5oFjjSSWMMFRwqdCfJg")).
                  setDefaultDirProvider(DEFAULT_DIR_PROVIDER).
                  build();
          Optional<ApiMessageAndVersion> expected = Optional.of(new ApiMessageAndVersion(
                  new PartitionChangeRecord().
                          setTopicId(FOO_ID).
                          setPartitionId(0).
                          setLeader(1).
                          setReplicas(List.of(3, 1, 5, 4)).
                          setDirectories(List.of(
                                  Uuid.fromString("fM5NKyWTQHqEihjIkUl99Q"),
                                  Uuid.fromString("iU2znv45Q9yQkOpkTSy3jA"),
                                  Uuid.fromString("RNJ5oFjjSSWMMFRwqdCfJg"),
                                  DirectoryId.UNASSIGNED
                          )),
                  (short) 1
          ));
          assertEquals(expected, built);
      }
  
      @Test
      public void testUpdateDirectories() {
          PartitionRegistration registration = new PartitionRegistration.Builder().
                  setReplicas(new int[]{2, 1, 3}).
                  setDirectories(new Uuid[]{
                          Uuid.fromString("S1zMYZczRjWmucidLqGA5g"),
                          Uuid.fromString("9eRNXTvFTsWUJObvW51V5A"),
                          Uuid.fromString("UpePYVBgRAi3c4ujQrf3Kg")
                  }).
                  setIsr(new int[]{2, 1, 3}).
                  setLeader(2).
                  setLeaderRecoveryState(LeaderRecoveryState.RECOVERED).
                  setLeaderEpoch(100).
                  setPartitionEpoch(200).
                  build();
          Optional<ApiMessageAndVersion> built = new PartitionChangeBuilder(registration, FOO_ID,
                  0, r -> true, MetadataVersion.latestTesting(), 2).
                  setDirectory(3, Uuid.fromString("pN1VKs9zRzK4APflpegAVg")).
                  setDirectory(1, DirectoryId.LOST).
                  setDefaultDirProvider(DEFAULT_DIR_PROVIDER).
                  build();
          Optional<ApiMessageAndVersion> expected = Optional.of(new ApiMessageAndVersion(
                  new PartitionChangeRecord().
                          setTopicId(FOO_ID).
                          setPartitionId(0).
                          setDirectories(List.of(
                                  Uuid.fromString("S1zMYZczRjWmucidLqGA5g"),
                                  DirectoryId.LOST,
                                  Uuid.fromString("pN1VKs9zRzK4APflpegAVg")
                          )),
                  (short) 2
          ));
          assertEquals(expected, built);
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void testEligibleLeaderReplicas_ElrCanBeElected(boolean lastKnownLeaderEnabled) {
          short version = 2;
          PartitionRegistration partition = new PartitionRegistration.Builder()
              .setReplicas(new int[] {1, 2, 3, 4})
              .setDirectories(DirectoryId.migratingArray(4))
              .setIsr(new int[] {1})
              .setElr(new int[] {3})
              .setLastKnownElr(lastKnownLeaderEnabled ? new int[] {} : new int[] {2})
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
              .setLeaderEpoch(100)
              .setPartitionEpoch(200)
              .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
  
          // Make replica 1 offline.
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> r != 1,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(isElrEnabled(version))
              .setUseLastKnownLeaderInBalancedRecovery(lastKnownLeaderEnabled)
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER);
  
          builder.setTargetIsr(List.of());
  
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(
              new PartitionChangeRecord()
                  .setTopicId(topicId)
                  .setPartitionId(0)
                  .setIsr(List.of(3))
                  .setEligibleLeaderReplicas(List.of(1))
                  .setLeader(3)
                  .setLeaderRecoveryState(LeaderRecoveryState.NO_CHANGE),
              version
          );
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          assertArrayEquals(new int[]{1}, partition.elr, partition.toString());
          assertArrayEquals(lastKnownLeaderEnabled ? new int[]{} : new int[]{2}, partition.lastKnownElr, partition.toString());
          assertArrayEquals(new int[]{3}, partition.isr, partition.toString());
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void testEligibleLeaderReplicas_IsrCanShrinkToZero(boolean lastKnownLeaderEnabled) {
          short version = 2;
          PartitionRegistration partition = new PartitionRegistration.Builder()
              .setReplicas(new int[] {1, 2, 3, 4})
              .setDirectories(new Uuid[]{
                      Uuid.fromString("MrTKKPEpRv66ZpWv4V7EBQ"),
                      Uuid.fromString("CkvgdEcWTVmdhfNuJXL0xA"),
                      Uuid.fromString("4a2coMsPRkSCsiTVWSksSw"),
                      Uuid.fromString("tmPdVjzASZ2ZqiS0cVJvtQ")
              })
              .setIsr(new int[] {1, 2, 3, 4})
              .setElr(new int[] {})
              .setLastKnownElr(new int[] {})
              .setLeader(1)
              .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
              .setLeaderEpoch(100)
              .setPartitionEpoch(200)
              .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
  
          // Mark all the replicas offline.
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> false,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(true)
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setUseLastKnownLeaderInBalancedRecovery(lastKnownLeaderEnabled);
  
          builder.setTargetIsr(List.of());
  
          PartitionChangeRecord record = new PartitionChangeRecord()
              .setTopicId(topicId)
              .setPartitionId(0)
              .setIsr(List.of())
              .setLeader(-1)
              .setLeaderRecoveryState(LeaderRecoveryState.NO_CHANGE)
              .setEligibleLeaderReplicas(List.of(1, 2, 3, 4));
  
          if (lastKnownLeaderEnabled) {
              record.setLastKnownElr(List.of(1));
          }
  
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(record, version);
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          assertArrayEquals(new int[]{1, 2, 3, 4}, partition.elr, partition.toString());
          if (lastKnownLeaderEnabled) {
              assertArrayEquals(new int[]{1}, partition.lastKnownElr, partition.toString());
              builder = new PartitionChangeBuilder(partition, topicId, 0, r -> false,
                      metadataVersionForPartitionChangeRecordVersion(version), 3)
                  .setElection(Election.PREFERRED)
                  .setEligibleLeaderReplicasEnabled(true)
                  .setUncleanShutdownReplicas(List.of(2))
                  .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
                  .setUseLastKnownLeaderInBalancedRecovery(lastKnownLeaderEnabled);
              PartitionChangeRecord changeRecord = (PartitionChangeRecord) builder.build().get().message();
              assertNull(changeRecord.lastKnownElr(), changeRecord.toString());
          } else {
              assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
          }
          assertArrayEquals(new int[]{}, partition.isr, partition.toString());
      }
  
      @Test
      public void testEligibleLeaderReplicas_ElectLastKnownLeader() {
          short version = 2;
          PartitionRegistration partition = new PartitionRegistration.Builder()
                  .setReplicas(new int[] {1, 2, 3, 4})
                  .setDirectories(DirectoryId.migratingArray(4))
                  .setIsr(new int[] {})
                  .setElr(new int[] {})
                  .setLastKnownElr(new int[] {1})
                  .setLeader(-1)
                  .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
                  .setLeaderEpoch(100)
                  .setPartitionEpoch(200)
                  .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
  
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> true,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setUseLastKnownLeaderInBalancedRecovery(true)
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setEligibleLeaderReplicasEnabled(true);
  
          builder.setTargetIsr(List.of());
  
          ApiMessageAndVersion expectedRecord = new ApiMessageAndVersion(
              new PartitionChangeRecord()
                  .setTopicId(topicId)
                  .setPartitionId(0)
                  .setIsr(List.of(1))
                  .setLeader(1)
                  .setLeaderRecoveryState(LeaderRecoveryState.RECOVERING.value())
                  .setLastKnownElr(List.of()),
              version
          );
          assertEquals(Optional.of(expectedRecord), builder.build());
          partition = partition.merge((PartitionChangeRecord) builder.build().get().message());
          assertArrayEquals(new int[]{}, partition.elr, partition.toString());
          assertArrayEquals(new int[]{}, partition.lastKnownElr, partition.toString());
          assertArrayEquals(new int[]{1}, partition.isr, partition.toString());
      }
  
      @Test
      public void testEligibleLeaderReplicas_ElectLastKnownLeaderShouldFail() {
          short version = 2;
          PartitionRegistration partition = new PartitionRegistration.Builder()
                  .setReplicas(new int[] {1, 2, 3, 4})
                  .setDirectories(new Uuid[]{
                          Uuid.fromString("zANDdMukTEqefOvHpmniMg"),
                          Uuid.fromString("Ui2Eq8rbRiuW7m7uiPTRyg"),
                          Uuid.fromString("MhgJOZrrTsKNcGM0XKK4aA"),
                          Uuid.fromString("Y25PaCAmRfyGIKxAThhBAw")
                  })
                  .setIsr(new int[] {})
                  .setElr(new int[] {3})
                  .setLastKnownElr(new int[] {1})
                  .setLeader(-1)
                  .setLeaderRecoveryState(LeaderRecoveryState.RECOVERED)
                  .setLeaderEpoch(100)
                  .setPartitionEpoch(200)
                  .build();
          Uuid topicId = Uuid.fromString("FbrrdcfiR-KC2CPSTHaJrg");
  
          PartitionChangeBuilder builder = new PartitionChangeBuilder(partition, topicId, 0, r -> r != 3,
                  metadataVersionForPartitionChangeRecordVersion(version), 3)
              .setElection(Election.PREFERRED)
              .setEligibleLeaderReplicasEnabled(true)
              .setDefaultDirProvider(DEFAULT_DIR_PROVIDER)
              .setUseLastKnownLeaderInBalancedRecovery(true);
  
          builder.setTargetIsr(List.of());
  
          // No change to the partition.
          assertEquals(Optional.empty(), builder.build());
      }
  }
  