Scala ZkClient类代码示例

//设置package包名称以及导入依赖的类
package com.groupon.dse.kafka.cluster.impl

import com.groupon.dse.configs.KafkaServerConfig
import com.groupon.dse.testutils.{EmbeddedKafka, TestDefaults}
import com.groupon.dse.zookeeper.ZkClientBuilder
import kafka.producer.{Producer, ProducerConfig}
import org.I0Itec.zkclient.ZkClient
import org.scalatest.{BeforeAndAfter, FlatSpec}

class ClusterImplTest extends FlatSpec with BeforeAndAfter {

  val kafkaTopic = TestDefaults.TestTopic
  val zkConnTimeout = 10000
  val zkSessionTimeout = 10000
  var producer: Producer[String, Array[Byte]] = _
  var embeddedKafka: EmbeddedKafka = _
  var cluster: ClusterImpl = _
  var zkConnect: String = _
  var kafkaServerConfigs: KafkaServerConfig = _
  var zkClient: ZkClient = _

  before {
    embeddedKafka = new EmbeddedKafka
    embeddedKafka.startCluster()
    producer = new Producer[String, Array[Byte]](new ProducerConfig(embeddedKafka.kafkaProducerProperties))
    zkConnect = embeddedKafka.zkServer.connectString
    kafkaServerConfigs = TestDefaults.testKafkaServerConfig(zkConnect)
    cluster = new ClusterImpl(kafkaServerConfigs)
    zkClient = ZkClientBuilder(zkConnect, zkConnTimeout, zkSessionTimeout)
  }

  after {
    zkClient.close()
    embeddedKafka.stopCluster()
  }

  "The topic list" must "have size 0 before producing" in {
    assert(cluster.topics(zkClient).size == 0)
  }

  "The topic list" must "have size 1 after producing" in {
    embeddedKafka.sendMessage(4, producer, kafkaTopic)
    assert(cluster.topics(zkClient).size == 1)
  }

  "The number of partitions for a topic" should "be 1 for 1 valid topic" in {
    embeddedKafka.sendMessage(4, producer, kafkaTopic)
    assert(cluster.partitions(List(kafkaTopic), zkClient).size == 1)
  }

  "The number of partitions" should "be 0 for an invalid topic" in {
    embeddedKafka.sendMessage(4, producer, kafkaTopic)
    assert(cluster.partitions(List("invalid_topic"), zkClient).size == 0)
  }

  "The number of partitions" should "be 1 for a valid and invalid topic" in {
    embeddedKafka.sendMessage(4, producer, kafkaTopic)
    assert(cluster.partitions(List(kafkaTopic, "invalid_topic"), zkClient).size == 1)
  }

} 

//设置package包名称以及导入依赖的类
package com.quantifind.kafka.core

import com.quantifind.utils.ZkUtilsWrapper
import org.I0Itec.zkclient.ZkClient
import org.apache.zookeeper.data.Stat
import org.mockito.Matchers._
import org.mockito.Mockito._
import org.mockito.{Matchers => MockitoMatchers, Mockito}
import org.scalatest._

class StormOffsetGetterSpec extends FlatSpec with ShouldMatchers {

  trait Fixture {

    val mockedZkClient = Mockito.mock(classOf[ZkClient])
    val zkOffsetBase = "/stormconsumers"
    val mockedZkUtil =  Mockito.mock(classOf[ZkUtilsWrapper])

    val offsetGetter = new StormOffsetGetter(mockedZkClient, zkOffsetBase, mockedZkUtil)
  }

  "StormOffsetGetter" should "be able to extract topic from persisted spout state" in new Fixture {

    val testGroup = "testgroup"
    val testTopic = "testtopic"
    val spoutState = s"""{
                        "broker": {
                            "host": "kafka.sample.net",
                            "port": 9092
                        },
                        "offset": 4285,
                        "partition": 1,
                        "topic": "${testTopic}",
                        "topology": {
                            "id": "fce905ff-25e0 -409e-bc3a-d855f 787d13b",
                            "name": "Test Topology"
                        }
                       }"""
    val ret = (spoutState, Mockito.mock(classOf[Stat]))
    when(mockedZkUtil.readData(MockitoMatchers.eq(mockedZkClient), anyString)).thenReturn(ret)

    val topics = offsetGetter.getTopicList(testGroup)
    
    topics.size shouldBe 1
    topics(0) shouldBe testTopic
  }
} 

//设置package包名称以及导入依赖的类
package com.quantifind.kafka.core

import com.quantifind.kafka.core.KafkaOffsetGetter.GroupTopicPartition
import com.quantifind.utils.ZkUtilsWrapper
import kafka.api.{OffsetRequest, OffsetResponse, PartitionOffsetsResponse}
import kafka.common.{OffsetAndMetadata, TopicAndPartition}
import kafka.consumer.SimpleConsumer
import org.I0Itec.zkclient.ZkClient
import org.mockito.Matchers._
import org.mockito.Mockito._
import org.mockito.{Matchers => MockitoMatchers, Mockito}
import org.scalatest._

class KafkaOffsetGetterSpec extends FlatSpec with ShouldMatchers {

  trait Fixture {

    val mockedZkClient = Mockito.mock(classOf[ZkClient])
    val mockedZkUtil =  Mockito.mock(classOf[ZkUtilsWrapper])
    val mockedConsumer = Mockito.mock(classOf[SimpleConsumer])
    val testPartitionLeader = 1

    val offsetGetter = new KafkaOffsetGetter(mockedZkClient, mockedZkUtil)
    offsetGetter.consumerMap += (testPartitionLeader -> Some(mockedConsumer))
  }

  "KafkaOffsetGetter" should "be able to build offset data for given partition" in new Fixture {

    val testGroup = "testgroup"
    val testTopic = "testtopic"
    val testPartition = 1

    val topicAndPartition = TopicAndPartition(testTopic, testPartition)
    val groupTopicPartition = GroupTopicPartition(testGroup, topicAndPartition)
    val offsetAndMetadata = OffsetAndMetadata(100, "meta", System.currentTimeMillis)
    KafkaOffsetGetter.offsetMap += (groupTopicPartition -> offsetAndMetadata)

    when(mockedZkUtil.getLeaderForPartition(MockitoMatchers.eq(mockedZkClient), MockitoMatchers.eq(testTopic), MockitoMatchers.eq(testPartition)))
      .thenReturn(Some(testPartitionLeader))

    val partitionErrorAndOffsets = Map(topicAndPartition -> PartitionOffsetsResponse(0,Seq(102)))
    val offsetResponse = OffsetResponse(1, partitionErrorAndOffsets)
    when(mockedConsumer.getOffsetsBefore(any[OffsetRequest])).thenReturn(offsetResponse)

    offsetGetter.processPartition(testGroup, testTopic, testPartition) match {
      case Some(offsetInfo) =>
        offsetInfo.topic shouldBe testTopic
        offsetInfo.group shouldBe testGroup
        offsetInfo.partition shouldBe testPartition
        offsetInfo.offset shouldBe 100
        offsetInfo.logSize shouldBe 102
      case None => fail("Failed to build offset data")
    }
    
  }
}