Embedded broker support different partition counts

Add `addTopics(NewTopic... topics)` to the embedded broker to enable
creation of topics with different partition counts.

* Polishing according PR comments

Author: garyrussell

spring-kafka-test/src/main/java/org/springframework/kafka/test/EmbeddedKafkaBroker.java

@@ -29,6 +29,7 @@
 import java.util.Map;
 import java.util.Properties;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.stream.Collectors;
 
@@ -87,6 +88,8 @@ public class EmbeddedKafkaBroker implements InitializingBean, DisposableBean {
 
 	public static final String SPRING_EMBEDDED_ZOOKEEPER_CONNECT = "spring.embedded.zookeeper.connect";
 
+	private static final int DEFAULT_ADMIN_TIMEOUT = 30;
+
 	private final int count;
 
 	private final boolean controlledShutdown;
@@ -107,6 +110,8 @@ public class EmbeddedKafkaBroker implements InitializingBean, DisposableBean {
 
 	private int[] kafkaPorts;
 
+	private int adminTimeout = DEFAULT_ADMIN_TIMEOUT;
+
 	public EmbeddedKafkaBroker(int count) {
 		this(count, false);
 	}
@@ -177,6 +182,16 @@ public EmbeddedKafkaBroker kafkaPorts(int... kafkaPorts) {
 		return this;
 	}
 
+	/**
+	 * Set the timeout in seconds for admin operations (e.g. topic creation, close).
+	 * Default 30 seconds.
+	 * @param adminTimeout the timeout.
+	 * @since 2.2
+	 */
+	public void setAdminTimeout(int adminTimeout) {
+		this.adminTimeout = adminTimeout;
+	}
+
 	@Override
 	public void afterPropertiesSet() {
 		this.zookeeper = new EmbeddedZookeeper();
@@ -218,33 +233,57 @@ private Properties createBrokerProperties(int i) {
 	}
 
 	/**
-	 * Create topics in the existing broker(s) using the configured number of partitions.
+	 * Add topics to the existing broker(s) using the configured number of partitions.
+	 * The broker(s) must be running.
 	 * @param topics the topics.
 	 */
-	private void createKafkaTopics(Set<String> topics) {
-		doWithAdmin(admin -> {
-			List<NewTopic> newTopics = topics.stream()
-					.map(t -> new NewTopic(t, this.partitionsPerTopic, (short) this.count))
-					.collect(Collectors.toList());
-			CreateTopicsResult createTopics = admin.createTopics(newTopics);
-			try {
-				createTopics.all().get();
-			}
-			catch (Exception e) {
-				throw new KafkaException(e);
-			}
-		});
+	public void addTopics(String... topics) {
+		Assert.notNull(this.zookeeper, "Broker must be started before this method can be called");
+		HashSet<String> set = new HashSet<>(Arrays.asList(topics));
+		createKafkaTopics(set);
+		this.topics.addAll(set);
 	}
 
+	/**
+	 * Add topics to the existing broker(s).
+	 * The broker(s) must be running.
+	 * @param topics the topics.
+	 * @since 2.2
+	 */
+	public void addTopics(NewTopic... topics) {
+		Assert.notNull(this.zookeeper, "Broker must be started before this method can be called");
+		for (NewTopic topic : topics) {
+			Assert.isTrue(this.topics.add(topic.name()), () -> "topic already exists: " + topic);
+			Assert.isTrue(topic.replicationFactor() <= this.count
+							&& (topic.replicasAssignments() == null
+							|| topic.replicasAssignments().size() <= this.count),
+					() -> "Embedded kafka does not support the requested replication factor: " + topic);
+		}
+
+		doWithAdmin(admin -> createTopics(admin, Arrays.asList(topics)));
+	}
 
 	/**
-	 * Add topics to the existing broker(s) using the configured number of partitions.
+	 * Create topics in the existing broker(s) using the configured number of partitions.
 	 * @param topics the topics.
 	 */
-	public void addTopics(String... topics) {
-		HashSet<String> set = new HashSet<>(Arrays.asList(topics));
-		createKafkaTopics(set);
-		this.topics.addAll(set);
+	private void createKafkaTopics(Set<String> topics) {
+		doWithAdmin(admin -> {
+			createTopics(admin,
+					topics.stream()
+							.map(t -> new NewTopic(t, this.partitionsPerTopic, (short) this.count))
+							.collect(Collectors.toList()));
+		});
+	}
+
+	private void createTopics(AdminClient admin, List<NewTopic> newTopics) {
+		CreateTopicsResult createTopics = admin.createTopics(newTopics);
+		try {
+			createTopics.all().get(this.adminTimeout, TimeUnit.SECONDS);
+		}
+		catch (Exception e) {
+			throw new KafkaException(e);
+		}
 	}
 
 	/**
@@ -254,9 +293,16 @@ public void addTopics(String... topics) {
 	public void doWithAdmin(java.util.function.Consumer<AdminClient> callback) {
 		Map<String, Object> adminConfigs = new HashMap<>();
 		adminConfigs.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, getBrokersAsString());
-		try (AdminClient admin = AdminClient.create(adminConfigs)) {
+		AdminClient admin = null;
+		try {
+			admin = AdminClient.create(adminConfigs);
 			callback.accept(admin);
 		}
+		finally {
+			if (admin != null) {
+				admin.close(this.adminTimeout, TimeUnit.SECONDS);
+			}
+		}
 	}
 
 	@Override

spring-kafka-test/src/main/java/org/springframework/kafka/test/context/EmbeddedKafka.java

@@ -53,6 +53,7 @@
  * @author Artem Bilan
  * @author Elliot Metsger
  * @author Zach Olauson
+ * @author Gary Russell
  *
  * @since 1.3
  *
@@ -87,8 +88,11 @@
 	int partitions() default 2;
 
 	/**
-	 * Topics that should be created
-	 * Topics may contain property placeholders, e.g. {@code topics = "${kafka.topic.one:topicOne}"}
+	 * Topics that should be created Topics may contain property placeholders, e.g.
+	 * {@code topics = "${kafka.topic.one:topicOne}"} The topics will be created with
+	 * {@link #partitions()} partitions; to provision other topics with other partition
+	 * counts call the {@code addTopics(NewTopic... topics)} method on the autowired
+	 * broker.
 	 * @return the topics to create
 	 */
 	String[] topics() default { };

spring-kafka-test/src/main/java/org/springframework/kafka/test/rule/KafkaEmbedded.java

@@ -131,15 +131,6 @@ public void afterPropertiesSet() throws Exception {
 		before();
 	}
 
-	/**
-	 * Add topics to the existing broker(s) using the configured number of partitions.
-	 * @param topics the topics.
-	 * @since 2.1.6
-	 */
-	public void addTopics(String... topics) {
-		getEmbeddedKafka().addTopics(topics);
-	}
-
 	/**
 	 * Create an {@link AdminClient}; invoke the callback and reliably close the
 	 * admin.

spring-kafka/src/main/java/org/springframework/kafka/listener/MessageListenerContainer.java

@@ -33,7 +33,7 @@
  * @author Gary Russell
  * @author Vladimir Tsanev
  */
-public interface MessageListenerContainer extends SmartLifecycle {
+	public interface MessageListenerContainer extends SmartLifecycle {
 
 	/**
 	 * Setup the message listener to use. Throws an {@link IllegalArgumentException}

spring-kafka/src/test/java/org/springframework/kafka/annotation/EnableKafkaIntegrationTests.java

@@ -17,6 +17,7 @@
 package org.springframework.kafka.annotation;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.BDDMockito.willAnswer;
@@ -40,6 +41,7 @@
 import javax.validation.ValidationException;
 import javax.validation.constraints.Max;
 
+import org.apache.kafka.clients.admin.NewTopic;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
@@ -686,6 +688,18 @@ public void testAddingTopics() {
 		int count = embeddedKafka.getTopics().size();
 		embeddedKafka.addTopics("testAddingTopics");
 		assertThat(embeddedKafka.getTopics().size()).isEqualTo(count + 1);
+		embeddedKafka.addTopics(new NewTopic("morePartitions", 10, (short) 1));
+		assertThat(embeddedKafka.getTopics().size()).isEqualTo(count + 2);
+		assertThatThrownBy(() -> embeddedKafka.addTopics(new NewTopic("morePartitions", 10, (short) 1)))
+			.isInstanceOf(IllegalArgumentException.class).hasMessageContaining("exists");
+		assertThatThrownBy(() -> embeddedKafka.addTopics(new NewTopic("morePartitions2", 10, (short) 2)))
+			.isInstanceOf(IllegalArgumentException.class).hasMessageContaining("replication");
+		Map<String, Object> consumerProps = new HashMap<>(this.consumerFactory.getConfigurationProperties());
+		consumerProps.put(ConsumerConfig.GROUP_ID_CONFIG, "testMultiReplying");
+		ConsumerFactory<Integer, String> cf = new DefaultKafkaConsumerFactory<>(consumerProps);
+		Consumer<Integer, String> consumer = cf.createConsumer();
+		assertThat(consumer.partitionsFor("morePartitions")).hasSize(10);
+		consumer.close();
 	}
 
 	@Test

src/reference/asciidoc/testing.adoc

@@ -103,41 +103,66 @@ Convenient constants `EmbeddedKafkaBroker.SPRING_EMBEDDED_KAFKA_BROKERS` and `Em
 With the `EmbeddedKafkaBroker.brokerProperties(Map<String, String>)` you can provide additional properties for the Kafka server(s).
 See https://kafka.apache.org/documentation/#brokerconfigs[Kafka Config] for more information about possible broker properties.
 
+==== Configuring Topics
+
+====
+[source, java]
+----
+public class MyTests {
+
+    @ClassRule
+    private static EmbeddedKafkaRule embeddedKafka = new EmbeddedKafkaRule(1, false, 5, "foo", "bar");
+
+    @Test
+    public void test() {
+        embeddedKafkaRule.getEmbeddedKafka()
+              .addTopics(new NewTopic("baz", 10, (short) 1), new NewTopic("qux", 15, (short) 1));
+        ...
+  	}
+
+}
+----
+====
+
+The above configuration will create topics `foo` and `bar` with 5 partitions, `baz` with 10 and `qux` with 15.
 
 ==== Using the Same Broker(s) for Multiple Test Classes
 
 There is no built-in support for this, but it can be achieved with something similar to the following:
 
+====
 [source, java]
 ----
 public final class EmbeddedKafkaHolder {
 
-	private static EmbeddedKafkaRule embeddedKafka = new EmbeddedKafkaRule(1, false);
+    private static EmbeddedKafkaRule embeddedKafka = new EmbeddedKafkaRule(1, false);
 
-	private static boolean started;
+    private static boolean started;
 
-	public static EmbeddedKafkaRule getEmbeddedKafka() {
-		if (!started) {
-			try {
-				embeddedKafka.before();
-			}
-			catch (Exception e) {
-				throw new KafkaException(e);
-			}
-			started = true;
-		}
-		return embeddedKafka;
-	}
+    public static EmbeddedKafkaRule getEmbeddedKafka() {
+        if (!started) {
+            try {
+                embeddedKafka.before();
+            }
+            catch (Exception e) {
+                throw new KafkaException(e);
+            }
+            started = true;
+        }
+        return embeddedKafka;
+    }
 
-	private EmbeddedKafkaHolder() {
-		super();
-	}
+    private EmbeddedKafkaHolder() {
+        super();
+    }
 
 }
 ----
+====
 
 And then, in each test class:
 
+====
 [source, java]
 ----
 static {
@@ -146,6 +171,7 @@ static {
 
 private static EmbeddedKafkaRule embeddedKafka = EmbeddedKafkaHolder.getEmbeddedKafka();
 ----
+====
 
 IMPORTANT: This example provides no mechanism for shutting down the broker(s) when all tests are complete.
 This could be a problem if, say, you run your tests in a Gradle daemon.
@@ -333,8 +359,8 @@ public class KafkaTemplateTests {
         final BlockingQueue<ConsumerRecord<Integer, String>> records = new LinkedBlockingQueue<>();
         container.setupMessageListener(new MessageListener<Integer, String>() {
 
-        	@Override
-        	public void onMessage(ConsumerRecord<Integer, String> record) {
+            @Override
+            public void onMessage(ConsumerRecord<Integer, String> record) {
                 System.out.println(record);
                 records.add(record);
             }