Kafka and zookeeper instances are now managed in a list, represented in case class instances. The start methods, which used to return Unit, now return references to the base type: EmbeddedServer. These references can be used to stop the specific servers using a new overload stop function.

Author: Thomas Heslin

embedded-kafka/src/main/scala/net/manub/embeddedkafka/EmbeddedKafka.scala

@@ -15,7 +15,6 @@ import org.apache.zookeeper.server.{ServerCnxnFactory, ZooKeeperServer}
 import org.scalatest.Suite
 
 import scala.collection.JavaConverters._
-import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 import scala.concurrent.duration._
 import scala.concurrent.{ExecutionContext, ExecutionContextExecutorService, TimeoutException}
@@ -29,81 +28,107 @@ trait EmbeddedKafka extends EmbeddedKafkaSupport {
 
 object EmbeddedKafka extends EmbeddedKafkaSupport {
 
-  private[this] var factory: Option[ServerCnxnFactory] = None
-  private[this] var broker: Option[KafkaServer] = None
-  private[this] val logsDirs = mutable.Buffer.empty[Directory]
+  private[this] var servers: Seq[EmbeddedServer] = Seq.empty
 
   /**
     * Starts a ZooKeeper instance and a Kafka broker in memory, using temporary directories for storing logs.
     * The log directories will be cleaned after calling the [[stop()]] method or on JVM exit, whichever happens earlier.
     *
     * @param config an implicit [[EmbeddedKafkaConfig]]
     */
-  def start()(implicit config: EmbeddedKafkaConfig): Unit = {
+  def start()(implicit config: EmbeddedKafkaConfig): EmbeddedK = {
     val zkLogsDir = Directory.makeTemp("zookeeper-logs")
     val kafkaLogsDir = Directory.makeTemp("kafka-logs")
 
-    factory = Option(startZooKeeper(config.zooKeeperPort, zkLogsDir))
-    broker = Option(startKafka(config, kafkaLogsDir))
+    val factory = EmbeddedZ(startZooKeeper(config.zooKeeperPort, zkLogsDir), zkLogsDir)
+    val broker = EmbeddedK(Option(factory), startKafka(config, kafkaLogsDir), kafkaLogsDir)
 
-    logsDirs ++= Seq(zkLogsDir, kafkaLogsDir)
+    servers :+= broker
+    broker
   }
 
   /**
     * Starts a Zookeeper instance in memory, storing logs in a specific location.
     *
     * @param zkLogsDir the path for the Zookeeper logs
     * @param config    an implicit [[EmbeddedKafkaConfig]]
+    * @return          an [[EmbeddedZ]] server
     */
   def startZooKeeper(zkLogsDir: Directory)(
-      implicit config: EmbeddedKafkaConfig): Unit = {
-    factory = Option(startZooKeeper(config.zooKeeperPort, zkLogsDir))
+    implicit config: EmbeddedKafkaConfig): EmbeddedZ = {
+    val factory = EmbeddedZ(startZooKeeper(config.zooKeeperPort, zkLogsDir), zkLogsDir)
+    servers :+= factory
+    factory
   }
 
   /**
     * Starts a Kafka broker in memory, storing logs in a specific location.
     *
-    * @param kafkaLogDir the path for the Kafka logs
-    * @param config      an implicit [[EmbeddedKafkaConfig]]
+    * @param kafkaLogsDir the path for the Kafka logs
+    * @param config       an implicit [[EmbeddedKafkaConfig]]
+    * @return             an [[EmbeddedK]] server
     */
-  def startKafka(kafkaLogDir: Directory)(
-      implicit config: EmbeddedKafkaConfig): Unit = {
-    broker = Option(startKafka(config, kafkaLogDir))
+  def startKafka(kafkaLogsDir: Directory)(
+    implicit config: EmbeddedKafkaConfig): EmbeddedK = {
+    val broker = EmbeddedK(startKafka(config, kafkaLogsDir), kafkaLogsDir)
+    servers :+= broker
+    broker
   }
 
   /**
-    * Stops the in memory ZooKeeper instance and Kafka broker, and deletes the log directories.
+    * Stops all in memory ZooKeeper instances and Kafka brokers, and deletes the log directories.
     */
   def stop(): Unit = {
-    stopKafka()
-    stopZooKeeper()
-    logsDirs.foreach(_.deleteRecursively())
-    logsDirs.clear()
+    servers.foreach(_.stop(true))
+    servers = Seq.empty
+  }
+
+  /**
+    * Stops a specific [[EmbeddedServer]] instance, and deletes the log directory.
+    *
+    * @param server the [[EmbeddedServer]] to be stopped.
+    */
+  def stop(server: EmbeddedServer): Unit = {
+    server.stop(true)
+    servers = servers.filter(x => x != server)
   }
 
   /**
-    * Stops the in memory Zookeeper instance, preserving the logs directory.
+    * Stops all in memory Zookeeper instances, preserving the logs directories.
     */
   def stopZooKeeper(): Unit = {
-    factory.foreach(_.shutdown())
-    factory = None
+    val factories = servers
+      .filter(_.isInstanceOf[EmbeddedZ])
+      .asInstanceOf[Seq[EmbeddedZ]]
+
+    factories
+      .foreach(_.stop(false))
+
+    servers = servers.filter(!factories.contains(_))
   }
 
   /**
-    * Stops the in memory Kafka instance, preserving the logs directory.
+    * Stops all in memory Kafka instances, preserving the logs directories.
     */
   def stopKafka(): Unit = {
-    broker.foreach { b =>
-      b.shutdown()
-      b.awaitShutdown()
-    }
-    broker = None
+    val brokers = servers
+      .filter(_.isInstanceOf[EmbeddedK])
+      .asInstanceOf[Seq[EmbeddedK]]
+
+    brokers
+      .foreach(_.stop(false))
+
+    servers = servers.filter(!brokers.contains(_))
   }
 
   /**
     * Returns whether the in memory Kafka and Zookeeper are running.
     */
-  def isRunning: Boolean = factory.nonEmpty && broker.nonEmpty
+  def isRunning: Boolean =
+    servers
+      .filter(_.isInstanceOf[EmbeddedK])
+      .asInstanceOf[Seq[EmbeddedK]]
+      .exists(_.factory.isDefined)
 }
 
 sealed trait EmbeddedKafkaSupport {

embedded-kafka/src/main/scala/net/manub/embeddedkafka/EmbeddedServer.scala

@@ -0,0 +1,42 @@
+package net.manub.embeddedkafka
+
+import kafka.server.KafkaServer
+import org.apache.zookeeper.server.ServerCnxnFactory
+
+import scala.reflect.io.Directory
+
+sealed trait EmbeddedServer {
+
+  def stop(clearLogs: Boolean): Unit
+}
+
+case class EmbeddedZ(factory: ServerCnxnFactory,
+                     logsDirs: Directory)(
+                      implicit config: EmbeddedKafkaConfig) extends EmbeddedServer {
+
+  override def stop(clearLogs: Boolean) = {
+    factory.shutdown()
+    if (clearLogs) logsDirs.deleteRecursively()
+  }
+}
+
+case class EmbeddedK(factory: Option[EmbeddedZ],
+                     broker: KafkaServer,
+                     logsDirs: Directory)(
+                      implicit config: EmbeddedKafkaConfig) extends EmbeddedServer {
+
+  override def stop(clearLogs: Boolean) = {
+    broker.shutdown()
+    broker.awaitShutdown()
+
+    factory.foreach(_.stop(clearLogs))
+
+    if (clearLogs) logsDirs.deleteRecursively()
+  }
+}
+
+object EmbeddedK {
+  def apply(broker: KafkaServer, logsDirs: Directory)(
+    implicit config: EmbeddedKafkaConfig): EmbeddedK =
+    EmbeddedK(None, broker, logsDirs)
+}

embedded-kafka/src/test/scala/net/manub/embeddedkafka/EmbeddedKafkaObjectSpec.scala

@@ -4,6 +4,7 @@ import org.apache.kafka.common.serialization.{StringDeserializer, StringSerializ
 import net.manub.embeddedkafka.EmbeddedKafka._
 
 import scala.collection.JavaConverters._
+import scala.reflect.io.Directory
 
 class EmbeddedKafkaObjectSpec extends EmbeddedKafkaSpecSupport {
 
@@ -34,12 +35,33 @@ class EmbeddedKafkaObjectSpec extends EmbeddedKafkaSpecSupport {
         EmbeddedKafka.stop()
       }
 
-      "multiple EmbeddedKafka can run in parallel" in {
+      "start and stop a specific Kafka" in {
+        val firstBroker = EmbeddedKafka.start()(EmbeddedKafkaConfig(kafkaPort = 7000, zooKeeperPort = 7001))
+        EmbeddedKafka.start()(EmbeddedKafkaConfig(kafkaPort = 8000, zooKeeperPort = 8001))
+
+        kafkaIsAvailable(7000)
+        zookeeperIsAvailable(7001)
+
+        kafkaIsAvailable(8000)
+        zookeeperIsAvailable(8001)
+
+        EmbeddedKafka.stop(firstBroker)
+
+        kafkaIsNotAvailable(7000)
+        zookeeperIsNotAvailable(7001)
+
+        kafkaIsAvailable(8000)
+        zookeeperIsAvailable(8001)
+
+        EmbeddedKafka.stop()
+      }
+
+      "start and stop multiple Kafka instances on specified ports" in {
         val someConfig = EmbeddedKafkaConfig(kafkaPort = 12345, zooKeeperPort = 32111)
-        EmbeddedKafka.start()(someConfig)
+        val someBroker = EmbeddedKafka.start()(someConfig)
 
         val someOtherConfig = EmbeddedKafkaConfig(kafkaPort = 23456, zooKeeperPort = 43211)
-        EmbeddedKafka.start()(someOtherConfig)
+        val someOtherBroker = EmbeddedKafka.start()(someOtherConfig)
 
         val topic = "publish_test_topic_1"
         val someMessage = "hello world!"
@@ -51,6 +73,8 @@ class EmbeddedKafkaObjectSpec extends EmbeddedKafkaSpecSupport {
         publishToKafka(topic, someMessage)(someConfig, serializer)
         publishToKafka(topic, someOtherMessage)(someOtherConfig, serializer)
 
+        // first
+
         val consumer = kafkaConsumer(someConfig, deserializer, deserializer)
         consumer.subscribe(List(topic).asJava)
 
@@ -60,6 +84,8 @@ class EmbeddedKafkaObjectSpec extends EmbeddedKafkaSpecSupport {
         val record = records.iterator().next
         record.value shouldBe someMessage
 
+        EmbeddedKafka.stop(someBroker)
+
         // second
 
         val anotherConsumer = kafkaConsumer(someOtherConfig, deserializer, deserializer)
@@ -71,18 +97,35 @@ class EmbeddedKafkaObjectSpec extends EmbeddedKafkaSpecSupport {
         val someOtherRecord = moreRecords.iterator().next
         someOtherRecord.value shouldBe someOtherMessage
 
-        EmbeddedKafka.stop()
-        EmbeddedKafka.stop()
+        EmbeddedKafka.stop(someOtherBroker)
       }
     }
 
-    "invoking the isRunnning method" should {
-      "return whether both Kafka and Zookeeper are running" in {
+    "invoking the isRunning method" should {
+      "return true when both Kafka and Zookeeper are running" in {
         EmbeddedKafka.start()
         EmbeddedKafka.isRunning shouldBe true
         EmbeddedKafka.stop()
         EmbeddedKafka.isRunning shouldBe false
       }
+
+      "return false when only Kafka is running" in {
+        val unmanagedZookeeper = EmbeddedKafka.startZooKeeper(6000, Directory.makeTemp("zookeeper-test-logs"))
+
+        EmbeddedKafka.startKafka(Directory.makeTemp("kafka-test-logs"))
+        EmbeddedKafka.isRunning shouldBe false
+        EmbeddedKafka.stop()
+        EmbeddedKafka.isRunning shouldBe false
+
+        unmanagedZookeeper.shutdown()
+      }
+
+      "return false when only Zookeeper is running" in {
+        EmbeddedKafka.startZooKeeper(Directory.makeTemp("zookeeper-test-logs"))
+        EmbeddedKafka.isRunning shouldBe false
+        EmbeddedKafka.stop()
+        EmbeddedKafka.isRunning shouldBe false
+      }
     }
   }
 }