From 92c77b329fc65fe37631e098c538bd9ec5a7155d Mon Sep 17 00:00:00 2001
From: Gary Russell <grussell@pivotal.io>
Date: Fri, 4 May 2018 16:27:00 -0400
Subject: [PATCH] GH-666: KafkaEmbedded doWithAdmin, addTopics

Resolves https://github.com/spring-projects/spring-kafka/issues/666

Allow arbitrary `AdminClient` operations and adding topics.

Polishing - PR Comments
---
 .../kafka/test/rule/KafkaEmbedded.java        | 45 +++++++++++++----
 src/reference/asciidoc/testing.adoc           | 48 +++++++++++++++++++
 2 files changed, 84 insertions(+), 9 deletions(-)

diff --git a/spring-kafka-test/src/main/java/org/springframework/kafka/test/rule/KafkaEmbedded.java b/spring-kafka-test/src/main/java/org/springframework/kafka/test/rule/KafkaEmbedded.java
index 33127c27..1ef47739 100644
--- a/spring-kafka-test/src/main/java/org/springframework/kafka/test/rule/KafkaEmbedded.java
+++ b/spring-kafka-test/src/main/java/org/springframework/kafka/test/rule/KafkaEmbedded.java
@@ -60,6 +60,7 @@ import org.springframework.retry.policy.SimpleRetryPolicy;
 import org.springframework.retry.support.RetryTemplate;
 import org.springframework.util.Assert;
 
+import kafka.common.KafkaException;
 import kafka.server.KafkaConfig;
 import kafka.server.KafkaServer;
 import kafka.server.NotRunning;
@@ -233,19 +234,45 @@ public class KafkaEmbedded extends ExternalResource implements KafkaRule, Initia
 				this.kafkaPorts[i] = TestUtils.boundPort(server, SecurityProtocol.PLAINTEXT);
 			}
 		}
-		Map<String, Object> adminConfigs = new HashMap<>();
-		adminConfigs.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, getBrokersAsString());
-		AdminClient admin = AdminClient.create(adminConfigs);
-		List<NewTopic> newTopics = Arrays.stream(this.topics)
-				.map(t -> new NewTopic(t, this.partitionsPerTopic, (short) this.count))
-				.collect(Collectors.toList());
-		CreateTopicsResult createTopics = admin.createTopics(newTopics);
-		createTopics.all().get();
-		admin.close();
+		addTopics(this.topics);
 		System.setProperty(SPRING_EMBEDDED_KAFKA_BROKERS, getBrokersAsString());
 		System.setProperty(SPRING_EMBEDDED_ZOOKEEPER_CONNECT, getZookeeperConnectionString());
 	}
 
+	/**
+	 * Add topics to the existing broker(s) using the configured number of partitions.
+	 * @param topics the topics.
+	 * @since 2.1
+	 */
+	public void addTopics(String... topics) {
+		doWithAdmin(admin -> {
+			List<NewTopic> newTopics = Arrays.stream(topics)
+					.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);
+			}
+		});
+	}
+
+	/**
+	 * Create an {@link AdminClient} invoke the callback and reliable close the
+	 * admin.
+	 * @param callback the callback.
+	 * @since 2.1
+	 */
+	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)) {
+			callback.accept(admin);
+		}
+	}
+
 	public Properties createBrokerProperties(int i) {
 		if (testUtilsCreateBrokerConfigMethod == null) {
 			return TestUtils.createBrokerConfig(i, this.zkConnect, this.controlledShutdown,
diff --git a/src/reference/asciidoc/testing.adoc b/src/reference/asciidoc/testing.adoc
index 6f114773..32292e2e 100644
--- a/src/reference/asciidoc/testing.adoc
+++ b/src/reference/asciidoc/testing.adoc
@@ -102,6 +102,54 @@ Convenient constants `KafkaEmbedded.SPRING_EMBEDDED_KAFKA_BROKERS` and `KafkaEmb
 With the `KafkaEmbedded.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.
 
+
+==== 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 KafkaEmbeddedHolder {
+
+	private static KafkaEmbedded kafkaEmbedded = new KafkaEmbedded(1, false);
+
+	private static boolean started;
+
+	public static KafkaEmbedded getKafkaEmbedded() {
+		if (!started) {
+			try {
+				kafkaEmbedded.before();
+			}
+			catch (Exception e) {
+				throw new KafkaException(e);
+			}
+			started = true;
+		}
+		return kafkaEmbedded;
+	}
+
+	private KafkaEmbeddedHolder() {
+		super();
+	}
+
+}
+----
+
+And then, in each test class:
+
+[source, java]
+----
+static {
+    KafkaEmbeddedHolder.getKafkaEmbedded().addTopics(topic1, topic2);
+}
+
+private static KafkaEmbedded embeddedKafka = KafkaEmbeddedHolder.getKafkaEmbedded();
+----
+
+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.
+You should not use this technique in such a situation, or use something to call `destroy()` on the `KafkaEmbedded` when your tests are complete.
+
 ==== @EmbeddedKafka Annotation
 It is generally recommended to use the rule as a `@ClassRule` to avoid starting/stopping the broker between tests (and use a different topic for each test).
 Starting with _version 2.0_, if you are using Spring's test application context caching, you can also declare a `KafkaEmbedded` bean, so a single broker can be used across multiple test classes.