Wait for partitions to be assigned when consuming from embedded topics

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2015-2017 the original author or authors.
+ * Copyright 2015-2018 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -72,6 +72,7 @@
  * @author Artem Bilan
  * @author Gary Russell
  * @author Kamill Sokol
+ * @author Elliot Kennedy
  */
 public class KafkaEmbedded extends ExternalResource implements KafkaRule, InitializingBean, DisposableBean {
 
@@ -370,23 +371,7 @@ public boolean isEmbedded() {
 	 * @throws Exception an exception.
 	 */
 	public void consumeFromAllEmbeddedTopics(Consumer<?, ?> consumer) throws Exception {
-		final CountDownLatch consumerLatch = new CountDownLatch(1);
-		consumer.subscribe(Arrays.asList(this.topics), new ConsumerRebalanceListener() {
-
-			@Override
-			public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
-			}
-
-			@Override
-			public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
-				consumerLatch.countDown();
-			}
-
-		});
-		consumer.poll(0); // force assignment
-		assertThat(consumerLatch.await(30, TimeUnit.SECONDS))
-				.as("Failed to be assigned partitions from the embedded topics")
-				.isTrue();
+		consumeFromEmbeddedTopics(consumer, this.topics);
 	}
 
 	/**
@@ -409,6 +394,7 @@ public void consumeFromEmbeddedTopics(Consumer<?, ?> consumer, String... topics)
 		for (String topic : topics) {
 			assertThat(this.topics).as("topic '" + topic + "' is not in embedded topic list").contains(topic);
 		}
+		final CountDownLatch consumerLatch = new CountDownLatch(1);
 		consumer.subscribe(Arrays.asList(topics), new ConsumerRebalanceListener() {
 
 			@Override
@@ -417,12 +403,17 @@ public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
 
 			@Override
 			public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
+				consumerLatch.countDown();
 				if (logger.isDebugEnabled()) {
 					logger.debug("partitions assigned: " + partitions);
 				}
 			}
 
 		});
+		consumer.poll(0); // force assignment
+		assertThat(consumerLatch.await(30, TimeUnit.SECONDS))
+				.as("Failed to be assigned partitions from the embedded topics")
+				.isTrue();
 		logger.debug("Subscription Initiated");
 	}
 

spring-kafka-test/src/test/java/org/springframework/kafka/test/hamcrest/AddressableEmbeddedBrokerTests.java → spring-kafka-test/src/test/java/org/springframework/kafka/test/rule/AddressableEmbeddedBrokerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017 the original author or authors.
+ * Copyright 2018 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.kafka.test.hamcrest;
+package org.springframework.kafka.test.rule;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -29,12 +29,12 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
-import org.springframework.kafka.test.rule.KafkaEmbedded;
 import org.springframework.test.context.junit4.SpringRunner;
 
 /**
  * @author Gary Russell
  * @author Kamill Sokol
+ * @author Elliot Kennedy
  * @since 1.3
  *
  */

spring-kafka/src/test/java/org/springframework/kafka/kstream/KafkaStreamsBranchTests.java

@@ -0,0 +1,179 @@
+/*
+ * Copyright 2018 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.kafka.kstream;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.serialization.Serdes;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.apache.kafka.streams.Consumed;
+import org.apache.kafka.streams.StreamsBuilder;
+import org.apache.kafka.streams.StreamsConfig;
+import org.apache.kafka.streams.kstream.KStream;
+import org.apache.kafka.streams.kstream.Produced;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.kafka.annotation.EnableKafka;
+import org.springframework.kafka.annotation.EnableKafkaStreams;
+import org.springframework.kafka.annotation.KafkaStreamsDefaultConfiguration;
+import org.springframework.kafka.core.DefaultKafkaConsumerFactory;
+import org.springframework.kafka.core.DefaultKafkaProducerFactory;
+import org.springframework.kafka.core.KafkaTemplate;
+import org.springframework.kafka.core.ProducerFactory;
+import org.springframework.kafka.test.context.EmbeddedKafka;
+import org.springframework.kafka.test.rule.KafkaEmbedded;
+import org.springframework.kafka.test.utils.KafkaTestUtils;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringRunner;
+
+/**
+ * @author Elliot Kennedy
+ */
+@RunWith(SpringRunner.class)
+@DirtiesContext
+@EmbeddedKafka(partitions = 1,
+		topics = {
+				KafkaStreamsBranchTests.TRUE_TOPIC,
+				KafkaStreamsBranchTests.FALSE_TOPIC,
+				KafkaStreamsBranchTests.TRUE_FALSE_INPUT_TOPIC },
+		brokerProperties = {
+				"auto.create.topics.enable=${topics.autoCreate:false}",
+				"delete.topic.enable=${topic.delete:true}"},
+		brokerPropertiesLocation = "classpath:/${broker.filename:broker}.properties")
+public class KafkaStreamsBranchTests {
+
+	public static final String TRUE_TOPIC = "true-output-topic";
+	public static final String FALSE_TOPIC = "false-output-topic";
+	public static final String TRUE_FALSE_INPUT_TOPIC = "input-topic";
+
+	@Autowired
+	private KafkaTemplate<String, String> kafkaTemplate;
+
+	@Autowired
+	private KafkaEmbedded kafkaEmbedded;
+
+	@Test
+	public void testBranchingStream() throws Exception {
+		Consumer<String, String> falseConsumer = createConsumer();
+		kafkaEmbedded.consumeFromAnEmbeddedTopic(falseConsumer, FALSE_TOPIC);
+
+		Consumer<String, String> trueConsumer = createConsumer();
+		kafkaEmbedded.consumeFromAnEmbeddedTopic(trueConsumer, TRUE_TOPIC);
+
+		kafkaTemplate.send(TRUE_FALSE_INPUT_TOPIC, UUID.randomUUID().toString(), String.valueOf(true));
+		kafkaTemplate.send(TRUE_FALSE_INPUT_TOPIC, UUID.randomUUID().toString(), String.valueOf(true));
+		kafkaTemplate.send(TRUE_FALSE_INPUT_TOPIC, UUID.randomUUID().toString(), String.valueOf(false));
+
+		ConsumerRecords<String, String> trueRecords = KafkaTestUtils.getRecords(trueConsumer);
+		ConsumerRecords<String, String> falseRecords = KafkaTestUtils.getRecords(falseConsumer);
+
+		List<String> trueValues = new ArrayList<>();
+		trueRecords.forEach(trueRecord -> trueValues.add(trueRecord.value()));
+
+		List<String> falseValues = new ArrayList<>();
+		falseRecords.forEach(falseRecord -> falseValues.add(falseRecord.value()));
+
+		assertThat(trueValues).containsExactly("true", "true");
+		assertThat(falseValues).containsExactly("false");
+	}
+
+	private Consumer<String, String> createConsumer() {
+		Map<String, Object> consumerProps = KafkaTestUtils.consumerProps(UUID.randomUUID().toString(), "false", kafkaEmbedded);
+		consumerProps.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 10000);
+
+		DefaultKafkaConsumerFactory<String, String> kafkaConsumerFactory = new DefaultKafkaConsumerFactory<>(consumerProps, new StringDeserializer(), new StringDeserializer());
+		return kafkaConsumerFactory.createConsumer();
+	}
+
+	@Configuration
+	@EnableKafka
+	@EnableKafkaStreams
+	public static class Config {
+
+		@Value("${" + KafkaEmbedded.SPRING_EMBEDDED_KAFKA_BROKERS + "}")
+		private String brokerAddresses;
+
+		@Bean
+		public ProducerFactory<Integer, String> producerFactory() {
+			return new DefaultKafkaProducerFactory<>(producerConfigs());
+		}
+
+		@Bean
+		public Map<String, Object> producerConfigs() {
+			return KafkaTestUtils.senderProps(this.brokerAddresses);
+		}
+
+		@Bean
+		public KafkaTemplate<Integer, String> template() {
+			KafkaTemplate<Integer, String> kafkaTemplate = new KafkaTemplate<>(producerFactory(), true);
+			kafkaTemplate.setDefaultTopic(TRUE_FALSE_INPUT_TOPIC);
+			return kafkaTemplate;
+		}
+
+		@Bean
+		public KafkaTemplate<String, String> kafkaTemplate() {
+			Map<String, Object> senderProperties = KafkaTestUtils.senderProps(this.brokerAddresses);
+			senderProperties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
+			senderProperties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
+
+			ProducerFactory<String, String> producerFactory = new DefaultKafkaProducerFactory<>(senderProperties);
+			return new KafkaTemplate<>(producerFactory);
+		}
+
+		@Bean(name = KafkaStreamsDefaultConfiguration.DEFAULT_STREAMS_CONFIG_BEAN_NAME)
+		public StreamsConfig kStreamsConfigs() {
+			Map<String, Object> props = new HashMap<>();
+			props.put(StreamsConfig.APPLICATION_ID_CONFIG, "testStreams");
+			props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, this.brokerAddresses);
+			props.put(StreamsConfig.DEFAULT_KEY_SERDE_CLASS_CONFIG, Serdes.String().getClass().getName());
+			props.put(StreamsConfig.DEFAULT_VALUE_SERDE_CLASS_CONFIG, Serdes.String().getClass().getName());
+			return new StreamsConfig(props);
+		}
+
+		@Bean
+		@SuppressWarnings("unchecked")
+		public KStream<String, String> trueFalseStream(StreamsBuilder streamsBuilder) {
+			KStream<String, String> trueFalseStream = streamsBuilder
+					.stream(TRUE_FALSE_INPUT_TOPIC, Consumed.with(Serdes.String(), Serdes.String()));
+
+			KStream<String, String>[] branches = trueFalseStream.branch((key, value) -> String.valueOf(true).equals(value), (key, value) -> String.valueOf(false).equals(value));
+
+			branches[0].to(TRUE_TOPIC, Produced.with(Serdes.String(), Serdes.String()));
+			branches[1].to(FALSE_TOPIC, Produced.with(Serdes.String(), Serdes.String()));
+
+			return trueFalseStream;
+		}
+
+	}
+
+}