/**
 * 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 kafka.zk

import kafka.security.authorizer.AclEntry.{WildcardHost, WildcardPrincipalString}
import kafka.server.{ConfigType, KafkaConfig}
import kafka.test.ClusterInstance
import kafka.test.annotation.{AutoStart, ClusterConfigProperty, ClusterTest, Type}
import kafka.test.junit.ClusterTestExtensions
import kafka.test.junit.ZkClusterInvocationContext.ZkClusterInstance
import kafka.testkit.{KafkaClusterTestKit, TestKitNodes}
import kafka.utils.{PasswordEncoder, TestUtils}
import org.apache.kafka.clients.admin._
import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
import org.apache.kafka.common.{TopicPartition, Uuid}
import org.apache.kafka.common.acl.AclOperation.{DESCRIBE, READ, WRITE}
import org.apache.kafka.common.acl.AclPermissionType.ALLOW
import org.apache.kafka.common.acl.{AccessControlEntry, AclBinding}
import org.apache.kafka.common.config.{ConfigResource, TopicConfig}
import org.apache.kafka.common.quota.{ClientQuotaAlteration, ClientQuotaEntity}
import org.apache.kafka.common.resource.PatternType.{LITERAL, PREFIXED}
import org.apache.kafka.common.resource.ResourcePattern
import org.apache.kafka.common.resource.ResourceType.TOPIC
import org.apache.kafka.common.security.auth.KafkaPrincipal
import org.apache.kafka.common.security.scram.internals.ScramCredentialUtils
import org.apache.kafka.common.serialization.StringSerializer
import org.apache.kafka.common.utils.SecurityUtils
import org.apache.kafka.image.{MetadataDelta, MetadataImage, MetadataProvenance}
import org.apache.kafka.metadata.authorizer.StandardAcl
import org.apache.kafka.metadata.migration.ZkMigrationLeadershipState
import org.apache.kafka.raft.RaftConfig
import org.apache.kafka.server.common.{ApiMessageAndVersion, MetadataVersion, ProducerIdsBlock}
import org.junit.jupiter.api.Assertions.{assertEquals, assertFalse, assertNotNull, assertTrue}
import org.junit.jupiter.api.Timeout
import org.junit.jupiter.api.extension.ExtendWith
import org.slf4j.LoggerFactory

import java.util
import java.util.concurrent.TimeUnit
import java.util.{Properties, UUID}
import scala.collection.Seq
import scala.jdk.CollectionConverters._


@ExtendWith(value = Array(classOf[ClusterTestExtensions]))
@Timeout(300)
class ZkMigrationIntegrationTest {

  val log = LoggerFactory.getLogger(classOf[ZkMigrationIntegrationTest])

  class MetadataDeltaVerifier {
    val metadataDelta = new MetadataDelta(MetadataImage.EMPTY)
    var offset = 0
    def accept(batch: java.util.List[ApiMessageAndVersion]): Unit = {
      batch.forEach(message => {
        metadataDelta.replay(message.message())
        offset += 1
      })
    }

    def verify(verifier: MetadataImage => Unit): Unit = {
      val image = metadataDelta.apply(new MetadataProvenance(offset, 0, 0))
      verifier.apply(image)
    }
  }

  @ClusterTest(
    brokers = 3, clusterType = Type.ZK, autoStart = AutoStart.YES,
    metadataVersion = MetadataVersion.IBP_3_4_IV0,
    serverProperties = Array(
      new ClusterConfigProperty(key="authorizer.class.name", value="kafka.security.authorizer.AclAuthorizer"),
      new ClusterConfigProperty(key="super.users", value="User:ANONYMOUS")
    )
  )
  def testMigrateAcls(clusterInstance: ClusterInstance): Unit = {
    val admin = clusterInstance.createAdminClient()

    val resource1 = new ResourcePattern(TOPIC, "foo-" + UUID.randomUUID(), LITERAL)
    val resource2 = new ResourcePattern(TOPIC, "bar-" + UUID.randomUUID(), LITERAL)
    val prefixedResource = new ResourcePattern(TOPIC, "bar-", PREFIXED)
    val username = "alice"
    val principal = new KafkaPrincipal(KafkaPrincipal.USER_TYPE, username)
    val wildcardPrincipal = SecurityUtils.parseKafkaPrincipal(WildcardPrincipalString)

    val acl1 = new AclBinding(resource1, new AccessControlEntry(principal.toString, WildcardHost, READ, ALLOW))
    val acl2 = new AclBinding(resource1, new AccessControlEntry(principal.toString, "192.168.0.1", WRITE, ALLOW))
    val acl3 = new AclBinding(resource2, new AccessControlEntry(principal.toString, WildcardHost, DESCRIBE, ALLOW))
    val acl4 = new AclBinding(prefixedResource, new AccessControlEntry(wildcardPrincipal.toString, WildcardHost, READ, ALLOW))

    val result = admin.createAcls(List(acl1, acl2, acl3, acl4).asJava)
    result.all().get

    val underlying = clusterInstance.asInstanceOf[ZkClusterInstance].getUnderlying()
    val zkClient = underlying.zkClient
    val migrationClient = ZkMigrationClient(zkClient, PasswordEncoder.noop())
    val verifier = new MetadataDeltaVerifier()
    migrationClient.readAllMetadata(batch => verifier.accept(batch), _ => { })
    verifier.verify { image =>
      val aclMap = image.acls().acls()
      assertEquals(4, aclMap.size())
      assertTrue(aclMap.values().containsAll(Seq(
        StandardAcl.fromAclBinding(acl1),
        StandardAcl.fromAclBinding(acl2),
        StandardAcl.fromAclBinding(acl3),
        StandardAcl.fromAclBinding(acl4)
      ).asJava))
    }
  }

  /**
   * Test ZkMigrationClient against a real ZooKeeper-backed Kafka cluster. This test creates a ZK cluster
   * and modifies data using AdminClient. The ZkMigrationClient is then used to read the metadata from ZK
   * as would happen during a migration. The generated records are then verified.
   */
  @ClusterTest(brokers = 3, clusterType = Type.ZK, metadataVersion = MetadataVersion.IBP_3_4_IV0)
  def testMigrate(clusterInstance: ClusterInstance): Unit = {
    val admin = clusterInstance.createAdminClient()
    val newTopics = new util.ArrayList[NewTopic]()
    newTopics.add(new NewTopic("test-topic-1", 2, 3.toShort)
      .configs(Map(TopicConfig.SEGMENT_BYTES_CONFIG -> "102400", TopicConfig.SEGMENT_MS_CONFIG -> "300000").asJava))
    newTopics.add(new NewTopic("test-topic-2", 1, 3.toShort))
    newTopics.add(new NewTopic("test-topic-3", 10, 3.toShort))
    val createTopicResult = admin.createTopics(newTopics)
    createTopicResult.all().get(60, TimeUnit.SECONDS)

    val quotas = new util.ArrayList[ClientQuotaAlteration]()
    quotas.add(new ClientQuotaAlteration(
      new ClientQuotaEntity(Map("user" -> "user1").asJava),
      List(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)).asJava))
    quotas.add(new ClientQuotaAlteration(
      new ClientQuotaEntity(Map("user" -> "user1", "client-id" -> "clientA").asJava),
      List(new ClientQuotaAlteration.Op("consumer_byte_rate", 800.0), new ClientQuotaAlteration.Op("producer_byte_rate", 100.0)).asJava))
    quotas.add(new ClientQuotaAlteration(
      new ClientQuotaEntity(Map("ip" -> "8.8.8.8").asJava),
      List(new ClientQuotaAlteration.Op("connection_creation_rate", 10.0)).asJava))
    admin.alterClientQuotas(quotas).all().get(60, TimeUnit.SECONDS)

    val zkClient = clusterInstance.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient
    val kafkaConfig = clusterInstance.asInstanceOf[ZkClusterInstance].getUnderlying.servers.head.config
    val zkConfigEncoder = kafkaConfig.passwordEncoderSecret match {
      case Some(secret) =>
        PasswordEncoder.encrypting(secret,
          kafkaConfig.passwordEncoderKeyFactoryAlgorithm,
          kafkaConfig.passwordEncoderCipherAlgorithm,
          kafkaConfig.passwordEncoderKeyLength,
          kafkaConfig.passwordEncoderIterations)
      case None => PasswordEncoder.noop()
    }

    val migrationClient = ZkMigrationClient(zkClient, zkConfigEncoder)
    var migrationState = migrationClient.getOrCreateMigrationRecoveryState(ZkMigrationLeadershipState.EMPTY)
    migrationState = migrationState.withNewKRaftController(3000, 42)
    migrationState = migrationClient.claimControllerLeadership(migrationState)

    val brokers = new java.util.HashSet[Integer]()
    val verifier = new MetadataDeltaVerifier()
    migrationClient.readAllMetadata(batch => verifier.accept(batch), brokerId => brokers.add(brokerId))
    assertEquals(Seq(0, 1, 2), brokers.asScala.toSeq)

    verifier.verify { image =>
      assertNotNull(image.topics().getTopic("test-topic-1"))
      assertEquals(2, image.topics().getTopic("test-topic-1").partitions().size())

      assertNotNull(image.topics().getTopic("test-topic-2"))
      assertEquals(1, image.topics().getTopic("test-topic-2").partitions().size())

      assertNotNull(image.topics().getTopic("test-topic-3"))
      assertEquals(10, image.topics().getTopic("test-topic-3").partitions().size())

      val clientQuotas = image.clientQuotas().entities()
      assertEquals(3, clientQuotas.size())
    }

    migrationState = migrationClient.releaseControllerLeadership(migrationState)
  }

  // SCRAM and Quota are intermixed. Test SCRAM Only here
  @ClusterTest(clusterType = Type.ZK, brokers = 3, metadataVersion = MetadataVersion.IBP_3_5_IV2, serverProperties = Array(
    new ClusterConfigProperty(key = "inter.broker.listener.name", value = "EXTERNAL"),
    new ClusterConfigProperty(key = "listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "advertised.listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "listener.security.protocol.map", value = "EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT"),
  ))
  def testDualWriteScram(zkCluster: ClusterInstance): Unit = {
    var admin = zkCluster.createAdminClient()
    createUserScramCredentials(admin).all().get(60, TimeUnit.SECONDS)
    admin.close()

    val zkClient = zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient

    // Bootstrap the ZK cluster ID into KRaft
    val clusterId = zkCluster.clusterId()
    val kraftCluster = new KafkaClusterTestKit.Builder(
      new TestKitNodes.Builder().
        setBootstrapMetadataVersion(MetadataVersion.IBP_3_5_IV2).
        setClusterId(Uuid.fromString(clusterId)).
        setNumBrokerNodes(0).
        setNumControllerNodes(1).build())
      .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
      .build()
    try {
      kraftCluster.format()
      kraftCluster.startup()
      val readyFuture = kraftCluster.controllers().values().asScala.head.controller.waitForReadyBrokers(3)

      // Enable migration configs and restart brokers
      log.info("Restart brokers in migration mode")
      val clientProps = kraftCluster.controllerClientProperties()
      val voters = clientProps.get(RaftConfig.QUORUM_VOTERS_CONFIG)
      zkCluster.config().serverProperties().put(KafkaConfig.MigrationEnabledProp, "true")
      zkCluster.config().serverProperties().put(RaftConfig.QUORUM_VOTERS_CONFIG, voters)
      zkCluster.config().serverProperties().put(KafkaConfig.ControllerListenerNamesProp, "CONTROLLER")
      zkCluster.config().serverProperties().put(KafkaConfig.ListenerSecurityProtocolMapProp, "CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT")
      zkCluster.rollingBrokerRestart()
      zkCluster.waitForReadyBrokers()
      readyFuture.get(30, TimeUnit.SECONDS)

      // Wait for migration to begin
      log.info("Waiting for ZK migration to begin")
      TestUtils.waitUntilTrue(() => zkClient.getControllerId.contains(3000), "Timed out waiting for KRaft controller to take over")

      // Alter the metadata
      log.info("Updating metadata with AdminClient")
      admin = zkCluster.createAdminClient()
      alterUserScramCredentials(admin).all().get(60, TimeUnit.SECONDS)

      // Verify the changes made to KRaft are seen in ZK
      log.info("Verifying metadata changes with ZK")
      verifyUserScramCredentials(zkClient)
    } finally {
      shutdownInSequence(zkCluster, kraftCluster)
    }
  }

  // SCRAM and Quota are intermixed. Test Quota Only here
  @ClusterTest(clusterType = Type.ZK, brokers = 3, metadataVersion = MetadataVersion.IBP_3_4_IV0, serverProperties = Array(
    new ClusterConfigProperty(key = "inter.broker.listener.name", value = "EXTERNAL"),
    new ClusterConfigProperty(key = "listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "advertised.listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "listener.security.protocol.map", value = "EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT"),
  ))
  def testDualWrite(zkCluster: ClusterInstance): Unit = {
    // Create a topic in ZK mode
    var admin = zkCluster.createAdminClient()
    val newTopics = new util.ArrayList[NewTopic]()
    newTopics.add(new NewTopic("test", 2, 3.toShort)
      .configs(Map(TopicConfig.SEGMENT_BYTES_CONFIG -> "102400", TopicConfig.SEGMENT_MS_CONFIG -> "300000").asJava))
    val createTopicResult = admin.createTopics(newTopics)
    createTopicResult.all().get(60, TimeUnit.SECONDS)
    admin.close()

    // Verify the configs exist in ZK
    val zkClient = zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient
    val propsBefore = zkClient.getEntityConfigs(ConfigType.Topic, "test")
    assertEquals("102400", propsBefore.getProperty(TopicConfig.SEGMENT_BYTES_CONFIG))
    assertEquals("300000", propsBefore.getProperty(TopicConfig.SEGMENT_MS_CONFIG))

    // Bootstrap the ZK cluster ID into KRaft
    val clusterId = zkCluster.clusterId()
    val kraftCluster = new KafkaClusterTestKit.Builder(
      new TestKitNodes.Builder().
        setBootstrapMetadataVersion(MetadataVersion.IBP_3_4_IV0).
        setClusterId(Uuid.fromString(clusterId)).
        setNumBrokerNodes(0).
        setNumControllerNodes(1).build())
      .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
      .build()
    try {
      kraftCluster.format()
      kraftCluster.startup()
      val readyFuture = kraftCluster.controllers().values().asScala.head.controller.waitForReadyBrokers(3)

      // Allocate a transactional producer ID while in ZK mode
      allocateProducerId(zkCluster.bootstrapServers())
      val producerIdBlock = readProducerIdBlock(zkClient)

      // Enable migration configs and restart brokers
      log.info("Restart brokers in migration mode")
      val clientProps = kraftCluster.controllerClientProperties()
      val voters = clientProps.get(RaftConfig.QUORUM_VOTERS_CONFIG)
      zkCluster.config().serverProperties().put(KafkaConfig.MigrationEnabledProp, "true")
      zkCluster.config().serverProperties().put(RaftConfig.QUORUM_VOTERS_CONFIG, voters)
      zkCluster.config().serverProperties().put(KafkaConfig.ControllerListenerNamesProp, "CONTROLLER")
      zkCluster.config().serverProperties().put(KafkaConfig.ListenerSecurityProtocolMapProp, "CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT")
      zkCluster.rollingBrokerRestart()
      zkCluster.waitForReadyBrokers()
      readyFuture.get(30, TimeUnit.SECONDS)

      // Wait for migration to begin
      log.info("Waiting for ZK migration to begin")
      TestUtils.waitUntilTrue(() => zkClient.getControllerId.contains(3000), "Timed out waiting for KRaft controller to take over")

      // Alter the metadata
      log.info("Updating metadata with AdminClient")
      admin = zkCluster.createAdminClient()
      alterTopicConfig(admin).all().get(60, TimeUnit.SECONDS)
      alterClientQuotas(admin).all().get(60, TimeUnit.SECONDS)

      // Verify the changes made to KRaft are seen in ZK
      log.info("Verifying metadata changes with ZK")
      verifyTopicConfigs(zkClient)
      verifyClientQuotas(zkClient)
      allocateProducerId(zkCluster.bootstrapServers())
      verifyProducerId(producerIdBlock, zkClient)

    } finally {
      shutdownInSequence(zkCluster, kraftCluster)
    }
  }

  // SCRAM and Quota are intermixed. Test both here
  @ClusterTest(clusterType = Type.ZK, brokers = 3, metadataVersion = MetadataVersion.IBP_3_5_IV2, serverProperties = Array(
    new ClusterConfigProperty(key = "inter.broker.listener.name", value = "EXTERNAL"),
    new ClusterConfigProperty(key = "listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "advertised.listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "listener.security.protocol.map", value = "EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT"),
  ))
  def testDualWriteQuotaAndScram(zkCluster: ClusterInstance): Unit = {
    var admin = zkCluster.createAdminClient()
    createUserScramCredentials(admin).all().get(60, TimeUnit.SECONDS)
    admin.close()

    val zkClient = zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient

    // Bootstrap the ZK cluster ID into KRaft
    val clusterId = zkCluster.clusterId()
    val kraftCluster = new KafkaClusterTestKit.Builder(
      new TestKitNodes.Builder().
        setBootstrapMetadataVersion(MetadataVersion.IBP_3_5_IV2).
        setClusterId(Uuid.fromString(clusterId)).
        setNumBrokerNodes(0).
        setNumControllerNodes(1).build())
      .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
      .build()
    try {
      kraftCluster.format()
      kraftCluster.startup()
      val readyFuture = kraftCluster.controllers().values().asScala.head.controller.waitForReadyBrokers(3)

      // Enable migration configs and restart brokers
      log.info("Restart brokers in migration mode")
      val clientProps = kraftCluster.controllerClientProperties()
      val voters = clientProps.get(RaftConfig.QUORUM_VOTERS_CONFIG)
      zkCluster.config().serverProperties().put(KafkaConfig.MigrationEnabledProp, "true")
      zkCluster.config().serverProperties().put(RaftConfig.QUORUM_VOTERS_CONFIG, voters)
      zkCluster.config().serverProperties().put(KafkaConfig.ControllerListenerNamesProp, "CONTROLLER")
      zkCluster.config().serverProperties().put(KafkaConfig.ListenerSecurityProtocolMapProp, "CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT")
      zkCluster.rollingBrokerRestart()
      zkCluster.waitForReadyBrokers()
      readyFuture.get(30, TimeUnit.SECONDS)

      // Wait for migration to begin
      log.info("Waiting for ZK migration to begin")
      TestUtils.waitUntilTrue(() => zkClient.getControllerId.contains(3000), "Timed out waiting for KRaft controller to take over")

      // Alter the metadata
      log.info("Updating metadata with AdminClient")
      admin = zkCluster.createAdminClient()
      alterUserScramCredentials(admin).all().get(60, TimeUnit.SECONDS)
      alterClientQuotas(admin).all().get(60, TimeUnit.SECONDS)

      // Verify the changes made to KRaft are seen in ZK
      log.info("Verifying metadata changes with ZK")
      verifyUserScramCredentials(zkClient)
      verifyClientQuotas(zkClient)
    } finally {
      shutdownInSequence(zkCluster, kraftCluster)
    }
  }

  @ClusterTest(clusterType = Type.ZK, brokers = 3, metadataVersion = MetadataVersion.IBP_3_4_IV0, serverProperties = Array(
    new ClusterConfigProperty(key = "inter.broker.listener.name", value = "EXTERNAL"),
    new ClusterConfigProperty(key = "listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "advertised.listeners", value = "PLAINTEXT://localhost:0,EXTERNAL://localhost:0"),
    new ClusterConfigProperty(key = "listener.security.protocol.map", value = "EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT"),
  ))
  def testNewAndChangedTopicsInDualWrite(zkCluster: ClusterInstance): Unit = {
    // Create a topic in ZK mode
    val topicName = "test"
    var admin = zkCluster.createAdminClient()
    val zkClient = zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying().zkClient

    // Bootstrap the ZK cluster ID into KRaft
    val clusterId = zkCluster.clusterId()
    val kraftCluster = new KafkaClusterTestKit.Builder(
      new TestKitNodes.Builder().
        setBootstrapMetadataVersion(MetadataVersion.IBP_3_4_IV0).
        setClusterId(Uuid.fromString(clusterId)).
        setNumBrokerNodes(0).
        setNumControllerNodes(1).build())
      .setConfigProp(KafkaConfig.MigrationEnabledProp, "true")
      .setConfigProp(KafkaConfig.ZkConnectProp, zkCluster.asInstanceOf[ZkClusterInstance].getUnderlying.zkConnect)
      .build()
    try {
      kraftCluster.format()
      kraftCluster.startup()
      val readyFuture = kraftCluster.controllers().values().asScala.head.controller.waitForReadyBrokers(3)

      // Enable migration configs and restart brokers
      log.info("Restart brokers in migration mode")
      val clientProps = kraftCluster.controllerClientProperties()
      val voters = clientProps.get(RaftConfig.QUORUM_VOTERS_CONFIG)
      zkCluster.config().serverProperties().put(KafkaConfig.MigrationEnabledProp, "true")
      zkCluster.config().serverProperties().put(RaftConfig.QUORUM_VOTERS_CONFIG, voters)
      zkCluster.config().serverProperties().put(KafkaConfig.ControllerListenerNamesProp, "CONTROLLER")
      zkCluster.config().serverProperties().put(KafkaConfig.ListenerSecurityProtocolMapProp, "CONTROLLER:PLAINTEXT,EXTERNAL:PLAINTEXT,PLAINTEXT:PLAINTEXT")
      zkCluster.rollingBrokerRestart()
      zkCluster.waitForReadyBrokers()
      readyFuture.get(30, TimeUnit.SECONDS)

      // Wait for migration to begin
      log.info("Waiting for ZK migration to begin")
      TestUtils.waitUntilTrue(() => zkClient.getControllerId.contains(3000), "Timed out waiting for KRaft controller to take over")

      // Alter the metadata
      log.info("Create new topic with AdminClient")
      admin = zkCluster.createAdminClient()
      val newTopics = new util.ArrayList[NewTopic]()
      newTopics.add(new NewTopic(topicName, 2, 3.toShort))
      val createTopicResult = admin.createTopics(newTopics)
      createTopicResult.all().get(60, TimeUnit.SECONDS)

      val existingPartitions = Seq(new TopicPartition(topicName, 0), new TopicPartition(topicName, 1))
      // Verify the changes made to KRaft are seen in ZK
      verifyTopicPartitionMetadata(topicName, existingPartitions, zkClient)

      log.info("Create new partitions with AdminClient")
      admin.createPartitions(Map(topicName -> NewPartitions.increaseTo(3)).asJava).all().get(60, TimeUnit.SECONDS)

      // Verify the changes seen in Zk.
      verifyTopicPartitionMetadata(topicName, existingPartitions ++ Seq(new TopicPartition(topicName, 2)), zkClient)
    } finally {
      zkCluster.stop()
      kraftCluster.close()
    }
  }

  def verifyTopicPartitionMetadata(topicName: String, partitions: Seq[TopicPartition], zkClient: KafkaZkClient): Unit = {
    val (topicIdReplicaAssignment, success) = TestUtils.computeUntilTrue(
      zkClient.getReplicaAssignmentAndTopicIdForTopics(Set(topicName)).headOption) {
      x => x.exists(_.assignment.size == partitions.size)
    }
    assertTrue(success, "Unable to find topic metadata in Zk")
    TestUtils.waitUntilTrue(() =>{
      val lisrMap = zkClient.getTopicPartitionStates(partitions.toSeq)
      lisrMap.size == partitions.size &&
        lisrMap.forall { case (tp, lisr) =>
          lisr.leaderAndIsr.leader >= 0 &&
            topicIdReplicaAssignment.exists(_.assignment(tp).replicas == lisr.leaderAndIsr.isr)
        }
    }, "Unable to find topic partition metadata")
  }

  def allocateProducerId(bootstrapServers: String): Unit = {
    val props = new Properties()
    props.put("bootstrap.servers", bootstrapServers)
    props.put("transactional.id", "some-transaction-id")
    val producer = new KafkaProducer[String, String](props, new StringSerializer(), new StringSerializer())
    producer.initTransactions()
    producer.beginTransaction()
    producer.send(new ProducerRecord[String, String]("test", "", "one"))
    producer.commitTransaction()
    producer.flush()
    producer.close()
  }

  def readProducerIdBlock(zkClient: KafkaZkClient): ProducerIdsBlock = {
    val (dataOpt, _) = zkClient.getDataAndVersion(ProducerIdBlockZNode.path)
    dataOpt.map(ProducerIdBlockZNode.parseProducerIdBlockData).get
  }

  def alterTopicConfig(admin: Admin): AlterConfigsResult = {
    val topicResource = new ConfigResource(ConfigResource.Type.TOPIC, "test")
    val alterConfigs = Seq(
      new AlterConfigOp(new ConfigEntry(TopicConfig.SEGMENT_BYTES_CONFIG, "204800"), AlterConfigOp.OpType.SET),
      new AlterConfigOp(new ConfigEntry(TopicConfig.SEGMENT_MS_CONFIG, null), AlterConfigOp.OpType.DELETE)
    ).asJavaCollection
    admin.incrementalAlterConfigs(Map(topicResource -> alterConfigs).asJava)
  }

  def alterClientQuotas(admin: Admin): AlterClientQuotasResult = {
    val quotas = new util.ArrayList[ClientQuotaAlteration]()
    quotas.add(new ClientQuotaAlteration(
      new ClientQuotaEntity(Map("user" -> "user1").asJava),
      List(new ClientQuotaAlteration.Op("consumer_byte_rate", 1000.0)).asJava))
    quotas.add(new ClientQuotaAlteration(
      new ClientQuotaEntity(Map("user" -> "user1", "client-id" -> "clientA").asJava),
      List(new ClientQuotaAlteration.Op("consumer_byte_rate", 800.0), new ClientQuotaAlteration.Op("producer_byte_rate", 100.0)).asJava))
    quotas.add(new ClientQuotaAlteration(
      new ClientQuotaEntity(Map("ip" -> "8.8.8.8").asJava),
      List(new ClientQuotaAlteration.Op("connection_creation_rate", 10.0)).asJava))
    admin.alterClientQuotas(quotas)
  }

  def createUserScramCredentials(admin: Admin): AlterUserScramCredentialsResult = {
    val alterations = new util.ArrayList[UserScramCredentialAlteration]()
    alterations.add(new UserScramCredentialUpsertion("user1",
        new ScramCredentialInfo(ScramMechanism.SCRAM_SHA_256, 8190), "password0"))
    admin.alterUserScramCredentials(alterations)
  }

  def alterUserScramCredentials(admin: Admin): AlterUserScramCredentialsResult = {
    val alterations = new util.ArrayList[UserScramCredentialAlteration]()
    alterations.add(new UserScramCredentialUpsertion("user1",
        new ScramCredentialInfo(ScramMechanism.SCRAM_SHA_256, 8191), "password1"))
    alterations.add(new UserScramCredentialUpsertion("user2",
        new ScramCredentialInfo(ScramMechanism.SCRAM_SHA_256, 8192), "password2"))
    admin.alterUserScramCredentials(alterations)
  }

  def verifyTopicConfigs(zkClient: KafkaZkClient): Unit = {
    TestUtils.retry(10000) {
      val propsAfter = zkClient.getEntityConfigs(ConfigType.Topic, "test")
      assertEquals("204800", propsAfter.getProperty(TopicConfig.SEGMENT_BYTES_CONFIG))
      assertFalse(propsAfter.containsKey(TopicConfig.SEGMENT_MS_CONFIG))
    }
  }

  def verifyClientQuotas(zkClient: KafkaZkClient): Unit = {
    TestUtils.retry(10000) {
      assertEquals("1000", zkClient.getEntityConfigs(ConfigType.User, "user1").getProperty("consumer_byte_rate"))
      assertEquals("800", zkClient.getEntityConfigs("users/user1/clients", "clientA").getProperty("consumer_byte_rate"))
      assertEquals("100", zkClient.getEntityConfigs("users/user1/clients", "clientA").getProperty("producer_byte_rate"))
      assertEquals("10", zkClient.getEntityConfigs(ConfigType.Ip, "8.8.8.8").getProperty("connection_creation_rate"))
    }
  }

  def verifyUserScramCredentials(zkClient: KafkaZkClient): Unit = {
    TestUtils.retry(10000) {
      val propertyValue1 = zkClient.getEntityConfigs(ConfigType.User, "user1").getProperty("SCRAM-SHA-256")
      val scramCredentials1 = ScramCredentialUtils.credentialFromString(propertyValue1)
      assertEquals(8191, scramCredentials1.iterations)

      val propertyValue2 = zkClient.getEntityConfigs(ConfigType.User, "user2").getProperty("SCRAM-SHA-256")
      assertNotNull(propertyValue2)
      val scramCredentials2 = ScramCredentialUtils.credentialFromString(propertyValue2)
      assertEquals(8192, scramCredentials2.iterations)
    }
  }

  def verifyProducerId(firstProducerIdBlock: ProducerIdsBlock, zkClient: KafkaZkClient): Unit = {
    TestUtils.retry(10000) {
      val producerIdBlock = readProducerIdBlock(zkClient)
      assertTrue(firstProducerIdBlock.firstProducerId() < producerIdBlock.firstProducerId())
    }
  }

  def shutdownInSequence(zkCluster: ClusterInstance, kraftCluster: KafkaClusterTestKit): Unit = {
    zkCluster.brokerIds().forEach(zkCluster.shutdownBroker(_))
    kraftCluster.close()
    zkCluster.stop()
  }
}
