|
| 1 | +/* |
| 2 | + * Licensed to the Apache Software Foundation (ASF) under one or more |
| 3 | + * contributor license agreements. See the NOTICE file distributed with |
| 4 | + * this work for additional information regarding copyright ownership. |
| 5 | + * The ASF licenses this file to You under the Apache License, Version 2.0 |
| 6 | + * (the "License"); you may not use this file except in compliance with |
| 7 | + * the License. You may obtain a copy of the License at |
| 8 | + * |
| 9 | + * http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | + * |
| 11 | + * Unless required by applicable law or agreed to in writing, software |
| 12 | + * distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | + * See the License for the specific language governing permissions and |
| 15 | + * limitations under the License. |
| 16 | + */ |
| 17 | + |
| 18 | +package org.apache.spark.rdd.kafka |
| 19 | + |
| 20 | +import scala.util.Random |
| 21 | + |
| 22 | +import kafka.serializer.StringDecoder |
| 23 | +import org.scalatest.BeforeAndAfter |
| 24 | + |
| 25 | +import org.apache.spark._ |
| 26 | +import org.apache.spark.SparkContext._ |
| 27 | +import org.apache.spark.streaming.kafka.KafkaStreamSuiteBase |
| 28 | + |
| 29 | +class KafkaRDDSuite extends KafkaStreamSuiteBase with BeforeAndAfter { |
| 30 | + var sc: SparkContext = _ |
| 31 | + before { |
| 32 | + setupKafka() |
| 33 | + } |
| 34 | + |
| 35 | + after { |
| 36 | + if (sc != null) { |
| 37 | + sc.stop |
| 38 | + sc = null |
| 39 | + } |
| 40 | + tearDownKafka() |
| 41 | + } |
| 42 | + |
| 43 | + test("Kafka RDD") { |
| 44 | + val sparkConf = new SparkConf().setMaster("local[4]").setAppName(this.getClass.getSimpleName) |
| 45 | + sc = new SparkContext(sparkConf) |
| 46 | + val topic = "topic1" |
| 47 | + val sent = Map("a" -> 5, "b" -> 3, "c" -> 10) |
| 48 | + createTopic(topic) |
| 49 | + produceAndSendMessage(topic, sent) |
| 50 | + |
| 51 | + val kafkaParams = Map("metadata.broker.list" -> s"localhost:$brokerPort", |
| 52 | + "group.id" -> s"test-consumer-${Random.nextInt(10000)}") |
| 53 | + |
| 54 | + val kc = new KafkaCluster(kafkaParams) |
| 55 | + |
| 56 | + val rdd = getRdd(kc, Set(topic)) |
| 57 | + assert(rdd.isDefined) |
| 58 | + assert(rdd.get.countByValue.size === sent.size) |
| 59 | + |
| 60 | + kc.setConsumerOffsets(kafkaParams("group.id"), rdd.get.untilOffsets) |
| 61 | + |
| 62 | + val rdd2 = getRdd(kc, Set(topic)) |
| 63 | + assert(rdd2.isDefined) |
| 64 | + assert(rdd2.get.count === 0) |
| 65 | + } |
| 66 | + |
| 67 | + private def getRdd(kc: KafkaCluster, topics: Set[String]) = { |
| 68 | + val groupId = kc.kafkaParams("group.id") |
| 69 | + for { |
| 70 | + topicPartitions <- kc.getPartitions(topics).right.toOption |
| 71 | + from <- kc.getConsumerOffsets(groupId, topicPartitions).right.toOption.orElse( |
| 72 | + kc.getEarliestLeaderOffsets(topicPartitions).right.toOption) |
| 73 | + until <- kc.getLatestLeaderOffsets(topicPartitions).right.toOption |
| 74 | + } yield { |
| 75 | + new KafkaRDD[String, String, StringDecoder, StringDecoder, String]( |
| 76 | + sc, kc.kafkaParams, from, until, mmd => mmd.message) |
| 77 | + } |
| 78 | + } |
| 79 | +} |
0 commit comments