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

Code before:
  override def setUp(testInfo: TestInfo): Unit = {
    super.setUp(testInfo)

    // create the test topic with all the brokers as replicas
    createTopic(topic, 2, brokerCount, adminClientConfig = this.adminClientConfig)
  }

  def awaitAssignment(consumer: Consumer[_, _], expectedAssignment: Set[TopicPartition])
  : Unit = {
    TestUtils.pollUntilTrue(consumer, () => consumer.assignment() == expectedAssignment.asJava,
      s"Timed out while awaiting expected assignment $expectedAssignment. " +
        s"The current assignment is ${consumer.assignment()}")
  }

  def awaitNonEmptyRecords[K, V](consumer: Consumer[K, V],
                                 partition: TopicPartition,
                                 pollTimeoutMs: Long = 100): ConsumerRecords[K, V] = {
    var result: ConsumerRecords[K, V] = null

    TestUtils.pollRecordsUntilTrue(consumer, (polledRecords: ConsumerRecords[K, V]) => {
      if (polledRecords.records(partition).asScala.nonEmpty) {
        result = polledRecords
        true
      } else {
        false
      }
    }, s"Consumer did not consume any messages for partition $partition before timeout.", JTestUtils.DEFAULT_MAX_WAIT_MS, pollTimeoutMs)

    result
  }

  /**
   * Creates 'consumerCount' consumers and consumer pollers, one per consumer; subscribes consumers to
   * 'topicsToSubscribe' topics, waits until consumers get topics assignment.
   *
   * When the function returns, consumer pollers will continue to poll until shutdown is called on every poller.
   *
   * @param consumerCount     number of consumers to create
   * @param topicsToSubscribe topics to which consumers will subscribe to
   * @param subscriptions     set of all topic partitions
   * @return collection of created consumers and collection of corresponding consumer pollers
   */
  def createConsumerGroupAndWaitForAssignment(consumerCount: Int,
                                              topicsToSubscribe: List[String],
                                              subscriptions: Set[TopicPartition]): (mutable.Buffer[Consumer[Array[Byte], Array[Byte]]], mutable.Buffer[ConsumerAssignmentPoller]) = {
    assertTrue(consumerCount <= subscriptions.size)
    val consumerGroup = mutable.Buffer[Consumer[Array[Byte], Array[Byte]]]()
    for (_ <- 0 until consumerCount)
      consumerGroup += createConsumer()

    // create consumer pollers, wait for assignment and validate it

Review comment (suggestion):
This part can be simplified.
```suggestion
      val hasRecords = !polledRecords.records(partition).isEmpty
      if (hasRecords) result = polledRecords
      hasRecords
```

Code after:
  override def setUp(testInfo: TestInfo): Unit = {
    super.setUp(testInfo)

    // create the test topic with all the brokers as replicas
    createTopic(topic, 2, brokerCount, adminClientConfig = this.adminClientConfig)
  }

  def awaitAssignment(consumer: Consumer[_, _], expectedAssignment: Set[TopicPartition])
  : Unit = {
    TestUtils.pollUntilTrue(consumer, () => consumer.assignment() == expectedAssignment.asJava,
      s"Timed out while awaiting expected assignment $expectedAssignment. " +
        s"The current assignment is ${consumer.assignment()}")
  }

  def awaitNonEmptyRecords[K, V](consumer: Consumer[K, V],
                                 partition: TopicPartition,
                                 pollTimeoutMs: Long = 100): ConsumerRecords[K, V] = {
    var result: ConsumerRecords[K, V] = null

    TestUtils.pollRecordsUntilTrue(consumer, (polledRecords: ConsumerRecords[K, V]) => {
      val hasRecords = !polledRecords.records(partition).isEmpty
      if (hasRecords) result = polledRecords
      hasRecords
    }, s"Consumer did not consume any messages for partition $partition before timeout.", JTestUtils.DEFAULT_MAX_WAIT_MS, pollTimeoutMs)
    result
  }

  /**
   * Creates 'consumerCount' consumers and consumer pollers, one per consumer; subscribes consumers to
   * 'topicsToSubscribe' topics, waits until consumers get topics assignment.
   *
   * When the function returns, consumer pollers will continue to poll until shutdown is called on every poller.
   *
   * @param consumerCount     number of consumers to create
   * @param topicsToSubscribe topics to which consumers will subscribe to
   * @param subscriptions     set of all topic partitions
   * @return collection of created consumers and collection of corresponding consumer pollers
   */
  def createConsumerGroupAndWaitForAssignment(consumerCount: Int,
                                              topicsToSubscribe: List[String],
                                              subscriptions: Set[TopicPartition]): (mutable.Buffer[Consumer[Array[Byte], Array[Byte]]], mutable.Buffer[ConsumerAssignmentPoller]) = {
    assertTrue(consumerCount <= subscriptions.size)
    val consumerGroup = mutable.Buffer[Consumer[Array[Byte], Array[Byte]]]()
    for (_ <- 0 until consumerCount)
      consumerGroup += createConsumer()

    // create consumer pollers, wait for assignment and validate it
    val consumerPollers = subscribeConsumers(consumerGroup, topicsToSubscribe)
    (consumerGroup, consumerPollers)
  }