Repository: apache/kafka
File: core/src/test/scala/integration/kafka/api/PlaintextAdminIntegrationTest.scala (Scala)

Code before:
        // Test listShareGroupOffsets with listShareGroupOffsetsSpec
        val groupSpecs = util.Map.of(testGroupId, new ListShareGroupOffsetsSpec().topicPartitions(util.List.of(tp1)))
        val parts = client.listShareGroupOffsets(groupSpecs).partitionsToOffsetAndMetadata(testGroupId).get()
        assertTrue(parts.containsKey(tp1))
        assertFalse(parts.containsKey(tp2))
      } finally {
        consumerThreads.foreach {
          case consumerThread =>
            consumerThread.interrupt()
            consumerThread.join()
        }
      }
    } finally {
      consumerSet.foreach(consumer => Utils.closeQuietly(consumer, "consumer"))
      Utils.closeQuietly(client, "adminClient")
    }
  }

  /**
   * Waits until the metadata for the given partition has fully propagated and become consistent across all brokers.
   *
   * @param partition The partition whose leader metadata should be verified across all brokers.
   */
  def waitForBrokerMetadataPropagation(partition: TopicPartition): Unit = {
    while (true) {
      val headBrokerLeaderId = brokers.head.metadataCache.getPartitionLeaderEndpoint(partition.topic, partition.partition(), listenerName).get.id()
      if (brokers.forall { broker =>
        val leaderId = broker.metadataCache.getPartitionLeaderEndpoint(partition.topic, partition.partition(), listenerName).get.id()
        leaderId == headBrokerLeaderId
      }) return
    }
  }

  @Test
  def testElectPreferredLeaders(): Unit = {
    client = createAdminClient

    val prefer0 = Seq(0, 1, 2)
    val prefer1 = Seq(1, 2, 0)
    val prefer2 = Seq(2, 0, 1)

    val partition1 = new TopicPartition("elect-preferred-leaders-topic-1", 0)
    createTopicWithAssignment(partition1.topic, Map[Int, Seq[Int]](partition1.partition -> prefer0))

    val partition2 = new TopicPartition("elect-preferred-leaders-topic-2", 0)
    createTopicWithAssignment(partition2.topic, Map[Int, Seq[Int]](partition2.partition -> prefer0))

    def preferredLeader(topicPartition: TopicPartition): Int = {
      val partitionMetadata = getTopicMetadata(client, topicPartition.topic).partitions.get(topicPartition.partition)
      val preferredLeaderMetadata = partitionMetadata.replicas.get(0)
      preferredLeaderMetadata.id

Review comment (performance):
It seems this helper method is used to ensure that "all" brokers see the same leader, right? If so, we could simplify the code
```scala
  def waitForBrokerMetadataPropagation(partition: TopicPartition): Unit = {
    while (brokers.map(_.metadataCache.getPartitionLeaderEndpoint(partition.topic, partition.partition(), listenerName))
      .filter(_.isPresent)
      .map(_.get())
      .toSet.size != 1)
      TimeUnit.MILLISECONDS.sleep(300)
  }
```

Code after:
        // Test listShareGroupOffsets with listShareGroupOffsetsSpec
        val groupSpecs = util.Map.of(testGroupId, new ListShareGroupOffsetsSpec().topicPartitions(util.List.of(tp1)))
        val parts = client.listShareGroupOffsets(groupSpecs).partitionsToOffsetAndMetadata(testGroupId).get()
        assertTrue(parts.containsKey(tp1))
        assertFalse(parts.containsKey(tp2))
      } finally {
        consumerThreads.foreach {
          case consumerThread =>
            consumerThread.interrupt()
            consumerThread.join()
        }
      }
    } finally {
      consumerSet.foreach(consumer => Utils.closeQuietly(consumer, "consumer"))
      Utils.closeQuietly(client, "adminClient")
    }
  }

  /**
   * Waits until the metadata for the given partition has fully propagated and become consistent across all brokers.
   *
   * @param partition The partition whose leader metadata should be verified across all brokers.
   */
    def waitForBrokerMetadataPropagation(partition: TopicPartition): Unit = {
      while (brokers.exists(_.metadataCache.getPartitionLeaderEndpoint(partition.topic, partition.partition(), listenerName).isEmpty) ||
        brokers.map(_.metadataCache.getPartitionLeaderEndpoint(partition.topic, partition.partition(), listenerName))
        .filter(_.isPresent)
        .map(_.get())
        .toSet.size != 1)
        TimeUnit.MILLISECONDS.sleep(300)
    }

  @Test
  def testElectPreferredLeaders(): Unit = {
    client = createAdminClient

    val prefer0 = Seq(0, 1, 2)
    val prefer1 = Seq(1, 2, 0)
    val prefer2 = Seq(2, 0, 1)

    val partition1 = new TopicPartition("elect-preferred-leaders-topic-1", 0)
    createTopicWithAssignment(partition1.topic, Map[Int, Seq[Int]](partition1.partition -> prefer0))

    val partition2 = new TopicPartition("elect-preferred-leaders-topic-2", 0)
    createTopicWithAssignment(partition2.topic, Map[Int, Seq[Int]](partition2.partition -> prefer0))

    def preferredLeader(topicPartition: TopicPartition): Int = {
      val partitionMetadata = getTopicMetadata(client, topicPartition.topic).partitions.get(topicPartition.partition)
      val preferredLeaderMetadata = partitionMetadata.replicas.get(0)
      preferredLeaderMetadata.id
    }