Merge pull request #24 from chris-zen/master

Added partitions and replicationFactor options to createCustomTopic

Author: manub

README.md

@@ -81,7 +81,7 @@ The `EmbeddedKafka` trait provides also some utility methods to interact with th
         
         def consumeFirstMessageFrom(topic: String): String
 
-        def createCustomTopic(topic: String, topicConfig: Map[String,String]): Unit
+        def createCustomTopic(topic: String, topicConfig: Map[String,String], partitions: Int, replicationFactor: Int): Unit
         
 ## Custom producers
 

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

@@ -242,15 +242,19 @@ sealed trait EmbeddedKafkaSupport {
   /**
     * Creates a topic with a custom configuration
     *
-    * @param topic       the topic name
-    * @param topicConfig per topic configuration [[Map]]
-    * @param config      an implicit [[EmbeddedKafkaConfig]]
+    * @param topic             the topic name
+    * @param topicConfig       per topic configuration [[Map]]
+    * @param partitions        number of partitions [[Int]]
+    * @param replicationFactor replication factor [[Int]]
+    * @param config            an implicit [[EmbeddedKafkaConfig]]
     */
-  def createCustomTopic(topic: String, topicConfig: Map[String,String] = Map.empty)(implicit config: EmbeddedKafkaConfig): Unit = {
+  def createCustomTopic(topic: String, topicConfig: Map[String,String] = Map.empty,
+    partitions: Int = 1, replicationFactor: Int = 1)(implicit config: EmbeddedKafkaConfig): Unit = {
+
     val zkUtils = ZkUtils(s"localhost:${config.zooKeeperPort}", zkSessionTimeoutMs, zkConnectionTimeoutMs, zkSecurityEnabled)
     val topicProperties = topicConfig.foldLeft(new Properties){case (props, (k,v)) => props.put(k,v); props}
 
-    try AdminUtils.createTopic(zkUtils, topic, 1, 1, topicProperties) finally zkUtils.close()
+    try AdminUtils.createTopic(zkUtils, topic, partitions, replicationFactor, topicProperties) finally zkUtils.close()
   }
 
 }

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

@@ -121,6 +121,24 @@ class EmbeddedKafkaSpec extends EmbeddedKafkaSpecSupport with EmbeddedKafka {
 
       }
     }
+
+    "create a topic with custom number of partitions" in {
+      implicit val config = EmbeddedKafkaConfig()
+      val topic = "test_custom_topic"
+
+      withRunningKafka {
+
+        createCustomTopic(topic, Map("cleanup.policy"->"compact"), partitions = 2)
+
+        val zkSessionTimeoutMs  = 10000
+        val zkConnectionTimeoutMs = 10000
+        val zkSecurityEnabled = false
+
+        val zkUtils = ZkUtils(s"localhost:${config.zooKeeperPort}", zkSessionTimeoutMs, zkConnectionTimeoutMs, zkSecurityEnabled)
+        try { AdminUtils.fetchTopicMetadataFromZk(topic, zkUtils).partitionsMetadata.size shouldBe 2 } finally zkUtils.close()
+
+      }
+    }
   }
 
   "the consumeFirstStringMessageFrom method" should {