Fixed test bug and refactored KafkaStreamSuite

Author: tdas

examples/scala-2.10/src/main/scala/org/apache/spark/examples/streaming/DirectKafkaWordCount.scala

@@ -26,7 +26,7 @@ import org.apache.spark.SparkConf
 /**
  * Consumes messages from one or more topics in Kafka and does wordcount.
  * Usage: DirectKafkaWordCount <brokers> <topics>
- *   <brokers> is a list of one or more zookeeper servers that make quorum
+ *   <brokers> is a list of one or more Kafka brokers
  *   <topics> is a list of one or more kafka topics to consume from
  *
  * Example:

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

@@ -171,7 +171,7 @@ object KafkaUtils {
     K: ClassTag,
     V: ClassTag,
     KD <: Decoder[K]: ClassTag,
-    VD <: Decoder[V]: ClassTag] (
+    VD <: Decoder[V]: ClassTag](
       sc: SparkContext,
       kafkaParams: Map[String, String],
       offsetRanges: Array[OffsetRange]
@@ -205,7 +205,7 @@ object KafkaUtils {
     V: ClassTag,
     KD <: Decoder[K]: ClassTag,
     VD <: Decoder[V]: ClassTag,
-    R: ClassTag] (
+    R: ClassTag](
       sc: SparkContext,
       kafkaParams: Map[String, String],
       offsetRanges: Array[OffsetRange],

external/kafka/src/test/java/org/apache/spark/streaming/kafka/JavaDirectKafkaStreamSuite.java

@@ -21,14 +21,11 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Random;
-import java.util.List;
-import java.util.ArrayList;
+import java.util.Arrays;
 
 import org.apache.spark.SparkConf;
-import org.apache.spark.streaming.Duration;
-import scala.Predef;
+
 import scala.Tuple2;
-import scala.collection.JavaConverters;
 
 import junit.framework.Assert;
 
@@ -74,12 +71,12 @@ public void testKafkaStream() throws InterruptedException {
     String topic1 = "topic1";
     String topic2 = "topic2";
 
-    List<String> topic1data = createTopicAndSendData(topic1);
-    List<String> topic2data = createTopicAndSendData(topic2);
+    String[] topic1data = createTopicAndSendData(topic1);
+    String[] topic2data = createTopicAndSendData(topic2);
 
     HashSet<String> sent = new HashSet<String>();
-    sent.addAll(topic1data);
-    sent.addAll(topic2data);
+    sent.addAll(Arrays.asList(topic1data));
+    sent.addAll(Arrays.asList(topic2data));
 
     HashMap<String, String> kafkaParams = new HashMap<String, String>();
     kafkaParams.put("metadata.broker.list", suiteBase.brokerAddress());
@@ -122,7 +119,7 @@ public String call(MessageAndMetadata<String, String> msgAndMd) throws Exception
 
     final HashSet<String> result = new HashSet<String>();
     unifiedStream.foreachRDD(
-        new Function<JavaRDD<String>, Void> () {
+        new Function<JavaRDD<String>, Void>() {
           @Override
           public Void call(org.apache.spark.api.java.JavaRDD<String> rdd) throws Exception {
             result.addAll(rdd.collect());
@@ -153,19 +150,10 @@ private HashMap<TopicAndPartition, Long> topicOffsetToMap(String topic, Long off
     return topicMap;
   }
 
-  private  List<String> createTopicAndSendData(String topic) {
-    List<String> data = java.util.Arrays.asList(topic+"-1", topic+"-2", topic+"-3");
-    HashMap<String, Integer> sent = new HashMap<String, Integer>();
-    for(String i: data) {
-      sent.put(i, 1);
-    }
-
+  private  String[] createTopicAndSendData(String topic) {
+    String[] data = { topic + "-1", topic + "-2", topic + "-3"};
     suiteBase.createTopic(topic);
-
-    HashMap<String, Object> tmp = new HashMap<String, Object>(sent);
-    suiteBase.produceAndSendMessage(topic,
-        JavaConverters.mapAsScalaMapConverter(tmp).asScala().toMap(
-            Predef.<Tuple2<String, Object>>conforms()));
+    suiteBase.sendMessages(topic, data);
     return data;
   }
 }

external/kafka/src/test/java/org/apache/spark/streaming/kafka/JavaKafkaStreamSuite.java

@@ -79,9 +79,10 @@ public void testKafkaStream() throws InterruptedException {
 
     suiteBase.createTopic(topic);
     HashMap<String, Object> tmp = new HashMap<String, Object>(sent);
-    suiteBase.produceAndSendMessage(topic,
+    suiteBase.sendMessages(topic,
         JavaConverters.mapAsScalaMapConverter(tmp).asScala().toMap(
-            Predef.<Tuple2<String, Object>>conforms()));
+            Predef.<Tuple2<String, Object>>conforms())
+    );
 
     HashMap<String, String> kafkaParams = new HashMap<String, String>();
     kafkaParams.put("zookeeper.connect", suiteBase.zkAddress());

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

@@ -27,7 +27,7 @@ import kafka.serializer.StringDecoder
 import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll}
 import org.scalatest.concurrent.{Eventually, Timeouts}
 
-import org.apache.spark.SparkConf
+import org.apache.spark.{SparkContext, SparkConf}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.streaming.{Milliseconds, StreamingContext, Time}
 import org.apache.spark.streaming.dstream.{DStream, InputDStream}
@@ -41,7 +41,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     .setMaster("local[4]")
     .setAppName(this.getClass.getSimpleName)
 
-
+  var sc: SparkContext = _
   var ssc: StreamingContext = _
   var testDir: File = _
 
@@ -56,18 +56,23 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
   after {
     if (ssc != null) {
       ssc.stop()
+      sc = null
+    }
+    if (sc != null) {
+      sc.stop()
     }
     if (testDir != null) {
       Utils.deleteRecursively(testDir)
     }
   }
 
-  test("basic receiving with multiple topics and smallest starting offset") {
-    val topics = Set("topic1", "topic2", "topic3")
+
+  test("basic stream receiving with multiple topics and smallest starting offset") {
+    val topics = Set("basic1", "basic2", "basic3")
     val data = Map("a" -> 7, "b" -> 9)
     topics.foreach { t =>
       createTopic(t)
-      produceAndSendMessage(t, data)
+      sendMessages(t, data)
     }
     val kafkaParams = Map(
       "metadata.broker.list" -> s"$brokerAddress",
@@ -107,6 +112,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     }
     ssc.stop()
   }
+
   test("receiving from largest starting offset") {
     val topic = "largest"
     val topicPartition = TopicAndPartition(topic, 0)
@@ -122,7 +128,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     }
 
     // Send some initial messages before starting context
-    produceAndSendMessage(topic, data)
+    sendMessages(topic, data)
     eventually(timeout(10 seconds), interval(20 milliseconds)) {
       assert(getLatestOffset() > 3)
     }
@@ -144,7 +150,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     stream.map { _._2 }.foreachRDD { rdd => collectedData ++= rdd.collect() }
     ssc.start()
     val newData = Map("b" -> 10)
-    produceAndSendMessage(topic, newData)
+    sendMessages(topic, newData)
     eventually(timeout(10 seconds), interval(50 milliseconds)) {
       collectedData.contains("b")
     }
@@ -167,7 +173,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     }
 
     // Send some initial messages before starting context
-    produceAndSendMessage(topic, data)
+    sendMessages(topic, data)
     eventually(timeout(10 seconds), interval(20 milliseconds)) {
       assert(getLatestOffset() >= 10)
     }
@@ -190,7 +196,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     stream.foreachRDD { rdd => collectedData ++= rdd.collect() }
     ssc.start()
     val newData = Map("b" -> 10)
-    produceAndSendMessage(topic, newData)
+    sendMessages(topic, newData)
     eventually(timeout(10 seconds), interval(50 milliseconds)) {
       collectedData.contains("b")
     }
@@ -211,7 +217,7 @@ class DirectKafkaStreamSuite extends KafkaStreamSuiteBase
     // Send data to Kafka and wait for it to be received
     def sendDataAndWaitForReceive(data: Seq[Int]) {
       val strings = data.map { _.toString}
-      produceAndSendMessage(topic, strings.map { _ -> 1}.toMap)
+      sendMessages(topic, strings.map { _ -> 1}.toMap)
       eventually(timeout(10 seconds), interval(50 milliseconds)) {
         assert(strings.forall { DirectKafkaStreamSuite.collectedData.contains })
       }

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

@@ -19,29 +19,29 @@ package org.apache.spark.streaming.kafka
 
 import scala.util.Random
 
-import org.scalatest.BeforeAndAfter
 import kafka.common.TopicAndPartition
+import org.scalatest.BeforeAndAfterAll
 
-class KafkaClusterSuite extends KafkaStreamSuiteBase with BeforeAndAfter {
-
-  val kafkaParams = Map("metadata.broker.list" -> s"$brokerAddress")
-  val kc = new KafkaCluster(kafkaParams)
+class KafkaClusterSuite extends KafkaStreamSuiteBase with BeforeAndAfterAll {
   val topic = "kcsuitetopic" + Random.nextInt(10000)
   val topicAndPartition = TopicAndPartition(topic, 0)
+  var kc: KafkaCluster = null
 
-  before {
+  override def beforeAll() {
     setupKafka()
     createTopic(topic)
-    produceAndSendMessage(topic, Map("a" -> 1))
+    sendMessages(topic, Map("a" -> 1))
+    kc = new KafkaCluster(Map("metadata.broker.list" -> s"$brokerAddress"))
   }
 
-  after {
+  override def afterAll() {
     tearDownKafka()
   }
 
   test("metadata apis") {
-    val leader = kc.findLeaders(Set(topicAndPartition)).right.get
-    assert(leader(topicAndPartition) === (brokerHost, brokerPort), "didn't get leader")
+    val leader = kc.findLeaders(Set(topicAndPartition)).right.get(topicAndPartition)
+    val leaderAddress = s"${leader._1}:${leader._2}"
+    assert(leaderAddress === brokerAddress, "didn't get leader")
 
     val parts = kc.getPartitions(Set(topic)).right.get
     assert(parts(topicAndPartition), "didn't get partitions")

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

@@ -46,9 +46,9 @@ class KafkaRDDSuite extends KafkaStreamSuiteBase with BeforeAndAfter {
     val topic = "topic1"
     val sent = Map("a" -> 5, "b" -> 3, "c" -> 10)
     createTopic(topic)
-    produceAndSendMessage(topic, sent)
+    sendMessages(topic, sent)
 
-    val kafkaParams = Map("metadata.broker.list" -> s"localhost:$brokerPort",
+    val kafkaParams = Map("metadata.broker.list" -> brokerAddress,
       "group.id" -> s"test-consumer-${Random.nextInt(10000)}")
 
     val kc = new KafkaCluster(kafkaParams)
@@ -65,14 +65,14 @@ class KafkaRDDSuite extends KafkaStreamSuiteBase with BeforeAndAfter {
 
     val rdd2 = getRdd(kc, Set(topic))
     val sent2 = Map("d" -> 1)
-    produceAndSendMessage(topic, sent2)
+    sendMessages(topic, sent2)
     // this is the "0 messages" case
     // make sure we dont get anything, since messages were sent after rdd was defined
     assert(rdd2.isDefined)
     assert(rdd2.get.count === 0)
 
     val rdd3 = getRdd(kc, Set(topic))
-    produceAndSendMessage(topic, Map("extra" -> 22))
+    sendMessages(topic, Map("extra" -> 22))
     // this is the "exactly 1 message" case
     // make sure we get exactly one message, despite there being lots more available
     assert(rdd3.isDefined)