Added a consumeFirstMessageFrom method.

Author: manub

src/main/scala/net/manub/embeddedkafka/EmbeddedKafka.scala

@@ -2,25 +2,34 @@ package net.manub.embeddedkafka
 
 import java.net.InetSocketAddress
 import java.util.Properties
+import java.util.concurrent.Executors
 
+import kafka.consumer.{Consumer, ConsumerConfig, Whitelist}
+import kafka.serializer.StringDecoder
 import kafka.server.{KafkaConfig, KafkaServer}
 import org.apache.kafka.clients.producer.{KafkaProducer, ProducerRecord}
 import org.apache.kafka.common.serialization.StringSerializer
 import org.apache.zookeeper.server.{ServerCnxnFactory, ZooKeeperServer}
 import org.scalatest.Suite
-import scala.collection.JavaConversions.mapAsJavaMap
 
+import scala.collection.JavaConversions.mapAsJavaMap
+import scala.concurrent.duration._
+import scala.concurrent.{Await, ExecutionContext, Future}
+import scala.language.postfixOps
 import scala.reflect.io.Directory
 
 trait EmbeddedKafka {
 
   this: Suite =>
 
+  val executorService = Executors.newFixedThreadPool(2)
+  implicit val executionContext = ExecutionContext.fromExecutorService(executorService)
+
   /**
    * Starts a ZooKeeper instance and a Kafka broker, then executes the body passed as a parameter.
    *
    * @param body the function to execute
-   * @param config an implicit {@see EmbeddedKafkaConfig}
+   * @param config an implicit [[EmbeddedKafkaConfig]]
    */
   def withRunningKafka(body: => Unit)(implicit config: EmbeddedKafkaConfig) = {
 
@@ -40,9 +49,9 @@ trait EmbeddedKafka {
    *
    * @param topic the topic to which publish the message (it will be auto-created)
    * @param message the message to publish
-   * @param config an implicit {@see EmbeddedKafkaConfig}
+   * @param config an implicit [[EmbeddedKafkaConfig]]
    */
-  def publishToKafka(topic: String, message: String)(implicit config: EmbeddedKafkaConfig) = {
+  def publishToKafka(topic: String, message: String)(implicit config: EmbeddedKafkaConfig): Unit = {
 
     val producerProps = Map(
       "bootstrap.servers" -> s"localhost:${config.kafkaPort}",
@@ -56,6 +65,37 @@ trait EmbeddedKafka {
     kafkaProducer.close()
   }
 
+
+  /**
+   * Consumes the first message available in a given topic, deserializing it as a String.
+   * Throws a [[java.util.concurrent.TimeoutException]] if a message is not available in 3 seconds.
+   *
+   * @param topic the topic to consume a message from
+   * @param config an implicit [[EmbeddedKafkaConfig]]
+   * @return the first message consumed from the given topic
+   */
+  def consumeFirstMessageFrom(topic: String)(implicit config: EmbeddedKafkaConfig): String = {
+    val props = new Properties()
+    props.put("group.id", "scalatest-embedded-kafka-spec")
+    props.put("zookeeper.connect", s"localhost:${config.zooKeeperPort}")
+    props.put("auto.offset.reset", "smallest")
+
+    val consumer = Consumer.create(new ConsumerConfig(props))
+
+    val filter = Whitelist(topic)
+    val messageStreams =
+      consumer.createMessageStreamsByFilter(filter, keyDecoder = new StringDecoder, valueDecoder = new StringDecoder)
+
+    val messageFuture = Future { messageStreams.head.iterator().next().message() }
+
+    try {
+      Await.result(messageFuture, 3 seconds)
+    } finally {
+      consumer.shutdown()
+    }
+  }
+
+
   private def startZooKeeper(zooKeeperPort: Int): ServerCnxnFactory = {
     val zkLogsDir = Directory.makeTemp("zookeeper-logs")
     val tickTime = 2000

src/test/scala/net/manub/embeddedkafka/EmbeddedKafkaSpec.scala

@@ -2,26 +2,30 @@ package net.manub.embeddedkafka
 
 import java.net.InetSocketAddress
 import java.util.Properties
+import java.util.concurrent.TimeoutException
 
 import akka.actor.{Actor, ActorRef, ActorSystem, Props}
 import akka.io.Tcp._
 import akka.io.{IO, Tcp}
 import akka.testkit.{ImplicitSender, TestKit}
-import kafka.consumer.{Consumer, ConsumerConfig, Whitelist}
+import kafka.consumer.{ConsumerConfig, Consumer, Whitelist}
 import kafka.serializer.StringDecoder
-import org.scalatest.concurrent.ScalaFutures
+import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
+import org.apache.kafka.common.serialization.StringSerializer
+import org.scalatest.concurrent.{JavaFutures, ScalaFutures}
 import org.scalatest.exceptions.TestFailedException
 import org.scalatest.time.{Milliseconds, Seconds, Span}
 import org.scalatest.{BeforeAndAfterAll, Matchers, WordSpecLike}
-import scala.concurrent.ExecutionContext.Implicits.global
 
+import scala.collection.JavaConversions._
+import scala.concurrent.ExecutionContext.Implicits.global
 import scala.concurrent.Future
 import scala.concurrent.duration._
 import scala.language.postfixOps
 
 class EmbeddedKafkaSpec
   extends TestKit(ActorSystem("embedded-kafka-spec")) with WordSpecLike with EmbeddedKafka with Matchers
-  with ImplicitSender with BeforeAndAfterAll with ScalaFutures {
+  with ImplicitSender with BeforeAndAfterAll with ScalaFutures with JavaFutures {
 
   implicit val config = PatienceConfig(Span(2, Seconds), Span(100, Milliseconds))
 
@@ -103,7 +107,7 @@ class EmbeddedKafkaSpec
 
   "the publishToKafka method" should {
 
-    "publish synchronously a message to Kafka as String" in {
+    "publishes asynchronously a message to Kafka as String" in {
 
       withRunningKafka {
 
@@ -134,6 +138,41 @@ class EmbeddedKafkaSpec
     }
   }
 
+  "the consumeFirstMessageFrom method" should {
+
+    "returns a message published to a topic" in {
+
+      withRunningKafka {
+
+        val message = "hello world!"
+        val topic = "test_topic"
+
+        val producer = new KafkaProducer[String, String](Map(
+          ProducerConfig.BOOTSTRAP_SERVERS_CONFIG -> s"localhost:6001",
+          ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG -> classOf[StringSerializer].getName,
+          ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG -> classOf[StringSerializer].getName
+        ))
+
+        whenReady(producer.send(new ProducerRecord[String, String](topic, message))) { _ =>
+          consumeFirstMessageFrom(topic) shouldBe message
+        }
+
+        producer.close()
+      }
+    }
+
+    "throws a TimeoutExeption when a message is not available" in {
+
+      withRunningKafka {
+
+        a[TimeoutException] shouldBe thrownBy {
+          consumeFirstMessageFrom("non_existing_topic")
+        }
+      }
+
+    }
+  }
+
   lazy val consumerConfigForEmbeddedKafka: ConsumerConfig = {
     val props = new Properties()
     props.put("group.id", "test")