org.apache.curator.test.TestingServer Scala Examples

package org.apache.openwhisk.core.invoker.test

import common.StreamLogging
import org.apache.curator.test.TestingServer
import org.junit.runner.RunWith
import org.scalatest.junit.JUnitRunner
import org.scalatest.{BeforeAndAfterEach, FlatSpec, Matchers}
import org.apache.openwhisk.core.invoker.InstanceIdAssigner

@RunWith(classOf[JUnitRunner])
class InstanceIdAssignerTests extends FlatSpec with Matchers with StreamLogging with BeforeAndAfterEach {
  behavior of "Id Assignment"

  private var zkServer: TestingServer = _

  override protected def beforeEach(): Unit = {
    zkServer = new TestingServer()
  }

  override protected def afterEach(): Unit = {
    zkServer.stop()
  }

  it should "assign fresh id" in {
    val assigner = new InstanceIdAssigner(zkServer.getConnectString)
    assigner.getId("foo") shouldBe 0
  }

  it should "reuse id if exists" in {
    val assigner = new InstanceIdAssigner(zkServer.getConnectString)
    assigner.getId("foo") shouldBe 0
    assigner.getId("bar") shouldBe 1
    assigner.getId("bar") shouldBe 1
  }

} 

package kafka.manager.utils

import org.apache.curator.framework.{CuratorFrameworkFactory, CuratorFramework}
import org.apache.curator.retry.ExponentialBackoffRetry
import org.apache.curator.test.TestingServer
import org.scalatest.{BeforeAndAfterAll, FunSuite}

import scala.reflect.ClassTag


trait CuratorAwareTest extends FunSuite with BeforeAndAfterAll with ZookeeperServerAwareTest {

  private[this] var curator: Option[CuratorFramework] = None

  override protected def beforeAll(): Unit = {
    super.beforeAll()
    val retryPolicy = new ExponentialBackoffRetry(1000, 3)
    val curatorFramework = CuratorFrameworkFactory.newClient(testServer.getConnectString, retryPolicy)
    curatorFramework.start
    curator = Some(curatorFramework)
  }

  override protected def afterAll(): Unit = {
    curator.foreach(_.close())
    super.afterAll()
  }

  protected def withCurator(fn: CuratorFramework => Unit): Unit = {
    curator.foreach(fn)
  }

  protected def produceWithCurator[T](fn: CuratorFramework => T) : T = {
    require(curator.isDefined,"Cannot produce with no curator defined!")
    fn(curator.get)
  }

  protected def checkError[T](fn: => Any)(implicit tag: ClassTag[T]): Unit = {
    try {
      fn
      throw new RuntimeException(s"expected ${tag.runtimeClass} , but no exceptions were thrown!")
    } catch {
      case UtilException(caught) =>
        if(!tag.runtimeClass.isAssignableFrom(caught.getClass)) {
          throw new RuntimeException(s"expected ${tag.runtimeClass} , found ${caught.getClass}, value=$caught")
        }
      case throwable: Throwable =>
        throw new RuntimeException(s"expected ${tag.runtimeClass} , found ${throwable.getClass}", throwable)
    }
  }

} 

package org.apache.kyuubi.ha.server

import java.io.File

import org.apache.curator.test.TestingServer

import org.apache.kyuubi.Logging
import org.apache.kyuubi.config.KyuubiConf
import org.apache.kyuubi.service.AbstractService


class EmbeddedZkServer private(name: String) extends AbstractService(name) with Logging {

  def this() = this(classOf[EmbeddedZkServer].getSimpleName)

  private var server: TestingServer = _

  override def initialize(conf: KyuubiConf): Unit = {
    this.conf = conf
    var port = conf.get(KyuubiConf.EMBEDDED_ZK_PORT)
    // When the client port is 0, the TestingServer will not randomly pick free local port to use
    // So adjust it to -1 to achieve what is common cognition.
    if (port == 0) port = -1
    val dataDir = conf.get(KyuubiConf.EMBEDDED_ZK_TEMP_DIR)
    server = new TestingServer(port, new File(dataDir), false)
    super.initialize(conf)
  }

  override def start(): Unit = {
    server.start()
    info(s"$getName is started at $getConnectString")
    super.start()
  }

  override def stop(): Unit = {
    if (server != null) {
      server.close()
    }
    server = null
    super.stop()
  }

  def getConnectString: String = {
    if (server == null) {
      null
    } else {
      server.getConnectString
    }
  }
} 

package yaooqinn.kyuubi.ha

import com.google.common.io.Files
import org.apache.curator.framework.{CuratorFramework, CuratorFrameworkFactory}
import org.apache.curator.retry.ExponentialBackoffRetry
import org.apache.curator.test.TestingServer
import org.apache.spark.{KyuubiConf, KyuubiSparkUtil, SparkConf, SparkFunSuite}
import org.apache.spark.KyuubiConf._

trait ZookeeperFunSuite extends SparkFunSuite{

  var zkServer: TestingServer = _
  var connectString: String = _
  val conf = new SparkConf(loadDefaults = true)
  KyuubiSparkUtil.setupCommonConfig(conf)
  conf.set(KyuubiConf.FRONTEND_BIND_PORT.key, "0")

  var zooKeeperClient: CuratorFramework = _

  override def beforeAll(): Unit = {
    zkServer = new TestingServer(2181, Files.createTempDir(), true)
    connectString = zkServer.getConnectString
    conf.set(HA_ZOOKEEPER_QUORUM.key, connectString)
    conf.set(HA_ZOOKEEPER_CONNECTION_BASESLEEPTIME.key, "100ms")
    conf.set(HA_ZOOKEEPER_SESSION_TIMEOUT.key, "15s")
    conf.set(HA_ZOOKEEPER_CONNECTION_MAX_RETRIES.key, "0")
    zooKeeperClient = CuratorFrameworkFactory.builder().connectString(connectString)
        .retryPolicy(new ExponentialBackoffRetry(1000, 3))
        .build()
    zooKeeperClient.start()
    super.beforeAll()
  }

  override def afterAll(): Unit = {
    Option(zooKeeperClient).foreach(_.close())
    Option(zkServer).foreach(_.stop())
    System.clearProperty(HA_ZOOKEEPER_QUORUM.key)
    System.clearProperty(HA_ENABLED.key)
    super.afterAll()
  }
}