[SPARK-2808][Streaming][Kafka] delay test until latest leader offset matches expected value

Author: koeninger

external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaTestUtils.scala

@@ -30,6 +30,7 @@ import scala.util.control.NonFatal
 
 import kafka.admin.AdminUtils
 import kafka.api.Request
+import kafka.common.TopicAndPartition
 import kafka.producer.{KeyedMessage, Producer, ProducerConfig}
 import kafka.serializer.StringEncoder
 import kafka.server.{KafkaConfig, KafkaServer}
@@ -228,6 +229,20 @@ private class KafkaTestUtils extends Logging {
     tryAgain(1)
   }
 
+  def waitUntilLeaderOffset(
+      kc: KafkaCluster,
+      topic: String,
+      partition: Int,
+      offset: Long): Unit = {
+    eventually(Time(10000), Time(100)) {
+      val tp = TopicAndPartition(topic, partition)
+      val llo = kc.getLatestLeaderOffsets(Set(tp)).right.get.apply(tp).offset
+      assert(
+        llo == offset,
+        s"$topic $partition $offset not reached after timeout")
+    }
+  }
+
   private def waitUntilMetadataIsPropagated(topic: String, partition: Int): Unit = {
     eventually(Time(10000), Time(100)) {
       assert(

external/kafka/src/test/scala/org/apache/spark/streaming/kafka/KafkaRDDSuite.scala

@@ -58,10 +58,12 @@ class KafkaRDDSuite extends FunSuite with BeforeAndAfterAll {
     val messages = Set("the", "quick", "brown", "fox")
     kafkaTestUtils.sendMessages(topic, messages.toArray)
 
-
     val kafkaParams = Map("metadata.broker.list" -> kafkaTestUtils.brokerAddress,
       "group.id" -> s"test-consumer-${Random.nextInt}")
 
+    val kc = new KafkaCluster(kafkaParams)
+    kafkaTestUtils.waitUntilLeaderOffset(kc, topic, 0, messages.size)
+
     val offsetRanges = Array(OffsetRange(topic, 0, 0, messages.size))
 
     val rdd =  KafkaUtils.createRDD[String, String, StringDecoder, StringDecoder](
@@ -84,14 +86,16 @@ class KafkaRDDSuite extends FunSuite with BeforeAndAfterAll {
 
     // this is the "lots of messages" case
     kafkaTestUtils.sendMessages(topic, sent)
+    val sentCount = sent.values.sum
+    kafkaTestUtils.waitUntilLeaderOffset(kc, topic, 0, sentCount)
+
     // rdd defined from leaders after sending messages, should get the number sent
     val rdd = getRdd(kc, Set(topic))
 
     assert(rdd.isDefined)
 
     val ranges = rdd.get.asInstanceOf[HasOffsetRanges].offsetRanges
     val rangeCount = ranges.map(o => o.untilOffset - o.fromOffset).sum
-    val sentCount = sent.values.sum
 
     assert(rangeCount === sentCount, "offset range didn't include all sent messages")
     assert(rdd.get.count === sentCount, "didn't get all sent messages")