Github user jerryshao commented on a diff in the pull request:
https://github.com/apache/spark/pull/5879#discussion_r29645169
--- Diff:
external/kafka/src/test/scala/org/apache/spark/streaming/kafka/DirectKafkaStreamSuite.scala
---
@@ -301,6 +301,49 @@ class DirectKafkaStreamSuite
ssc.stop()
}
+ test("Direct Kafka stream report input information") {
+ val topic = "report-test"
+ val data = Map("a" -> 7, "b" -> 9)
+ kafkaTestUtils.createTopic(topic)
+ kafkaTestUtils.sendMessages(topic, data)
+
+ val totalSent = data.values.sum
+ val kafkaParams = Map(
+ "metadata.broker.list" -> kafkaTestUtils.brokerAddress,
+ "auto.offset.reset" -> "smallest"
+ )
+
+ import DirectKafkaStreamSuite._
+ ssc = new StreamingContext(sparkConf, Milliseconds(200))
+ val collector = new InputInfoCollector
+ ssc.addStreamingListener(collector)
+
+ val stream = withClue("Error creating direct stream") {
+ KafkaUtils.createDirectStream[String, String, StringDecoder,
StringDecoder](
+ ssc, kafkaParams, Set(topic))
+ }
+
+ val allReceived = new ArrayBuffer[(String, String)]
+
+ stream.foreachRDD { rdd => allReceived ++= rdd.collect() }
+ ssc.start()
+ eventually(timeout(20000.milliseconds), interval(200.milliseconds)) {
+ assert(allReceived.size === totalSent,
+ "didn't get expected number of messages, messages:\n" +
allReceived.mkString("\n"))
+ }
+ ssc.stop()
+
+ // Calculate all the record number collected in the StreamingListener.
+ val numRecordsSubmitted =
collector.streamIdToNumRecordsSubmitted.map(_.values.sum).sum
--- End diff --
>There is probably race condition here. The collector may not have received
the batch completed signal when the allReceived.size == totalSent is satisfied
and the context is stopped. Better to put all of these asserts under the
eventually.
Not sure why all the assert should be in the `eventually`, from my
understanding it is OK the last signal is missed, since we only test the total
number of completed records.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
