From 600e88ec56c087ea25113771757e2b27399ca9b8 Mon Sep 17 00:00:00 2001
From: Cameron Mayfield <cameronjmayfield@gmail.com>
Date: Tue, 25 Jun 2019 07:12:37 -0700
Subject: [PATCH] GH-272: Add KafkaMDrivenChAdapterSpec.payloadType

Fixes https://github.com/spring-projects/spring-integration-kafka/issues/272

Add `payloadType` option into `KafkaMessageDrivenChannelAdapterSpec`
---
 .../KafkaMessageDrivenChannelAdapterSpec.java | 13 +++++
 .../inbound/MessageDrivenAdapterTests.java    | 47 ++++++++++++++++---
 2 files changed, 53 insertions(+), 7 deletions(-)

diff --git a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaMessageDrivenChannelAdapterSpec.java b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaMessageDrivenChannelAdapterSpec.java
index 4a03b9e990..949951be76 100644
--- a/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaMessageDrivenChannelAdapterSpec.java
+++ b/spring-integration-kafka/src/main/java/org/springframework/integration/kafka/dsl/KafkaMessageDrivenChannelAdapterSpec.java
@@ -46,6 +46,7 @@ import org.springframework.util.Assert;
  *
  * @author Artem Bilan
  * @author Gary Russell
+ * @author Cameron Mayfield
  *
  * @since 3.0
  */
@@ -141,6 +142,18 @@ public class KafkaMessageDrivenChannelAdapterSpec<K, V, S extends KafkaMessageDr
 		return _this();
 	}
 
+	/**
+	 * When using a type-aware message converter (such as {@code StringJsonMessageConverter},
+	 * set the payload type the converter should create. Defaults to {@link Object}.
+	 * @param payloadType the type.
+	 * @return the spec
+	 * @since 3.2.0
+	 */
+	public S payloadType(Class payloadType) {
+		this.target.setPayloadType(payloadType);
+		return _this();
+	}
+
 	/**
 	 * The {@code boolean} flag to specify the order how
 	 * {@link org.springframework.kafka.listener.adapter.RetryingMessageListenerAdapter}
diff --git a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java
index 99908322a9..9a35e3a327 100644
--- a/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java
+++ b/spring-integration-kafka/src/test/java/org/springframework/integration/kafka/inbound/MessageDrivenAdapterTests.java
@@ -54,6 +54,7 @@ import org.springframework.integration.StaticMessageHeaderAccessor;
 import org.springframework.integration.channel.DirectChannel;
 import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.handler.advice.ErrorMessageSendingRecoverer;
+import org.springframework.integration.kafka.dsl.Kafka;
 import org.springframework.integration.kafka.inbound.KafkaMessageDrivenChannelAdapter.ListenerMode;
 import org.springframework.integration.kafka.support.RawRecordHeaderErrorMessageStrategy;
 import org.springframework.integration.support.MessageBuilder;
@@ -92,6 +93,7 @@ import org.springframework.retry.support.RetryTemplate;
  * @author Gary Russell
  * @author Artem Bilan
  * @author Biju Kunjummen
+ * @author Cameron Mayfield
  *
  * @since 2.0
  *
@@ -108,9 +110,11 @@ public class MessageDrivenAdapterTests {
 
 	private static String topic5 = "testTopic5";
 
+	private static String topic6 = "testTopic6";
+
 	@ClassRule
 	public static EmbeddedKafkaRule embeddedKafkaRule =
-			new EmbeddedKafkaRule(1, true, topic1, topic2, topic3, topic4, topic5);
+			new EmbeddedKafkaRule(1, true, topic1, topic2, topic3, topic4, topic5,  topic6);
 
 	private static EmbeddedKafkaBroker embeddedKafka = embeddedKafkaRule.getEmbeddedKafka();
 
@@ -390,7 +394,7 @@ public class MessageDrivenAdapterTests {
 	}
 
 	@Test
-	public void testInboundJson() throws Exception {
+	public void testInboundJson() {
 		Map<String, Object> props = KafkaTestUtils.consumerProps("test3", "true", embeddedKafka);
 		props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
 		DefaultKafkaConsumerFactory<Integer, String> cf = new DefaultKafkaConsumerFactory<>(props);
@@ -428,17 +432,46 @@ public class MessageDrivenAdapterTests {
 		assertThat(headers.get("foo")).isEqualTo("bar");
 		assertThat(received.getPayload()).isInstanceOf(Map.class);
 
-		adapter.setPayloadType(Foo.class);
+		adapter.stop();
+	}
+
+	@Test
+	public void testInboundJsonWithPayload() {
+		Map<String, Object> props = KafkaTestUtils.consumerProps("test6", "true", embeddedKafka);
+		props.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");
+		DefaultKafkaConsumerFactory<Integer, Foo> cf = new DefaultKafkaConsumerFactory<>(props);
+		ContainerProperties containerProps = new ContainerProperties(topic6);
+		KafkaMessageListenerContainer<Integer, Foo> container =
+				new KafkaMessageListenerContainer<>(cf, containerProps);
+
+		KafkaMessageDrivenChannelAdapter<Integer, Foo> adapter = Kafka.messageDrivenChannelAdapter(container, ListenerMode.record)
+				.recordMessageConverter(new StringJsonMessageConverter())
+				.payloadType(Foo.class)
+				.get();
+		QueueChannel out = new QueueChannel();
+		adapter.setOutputChannel(out);
+		adapter.afterPropertiesSet();
+		adapter.start();
+		ContainerTestUtils.waitForAssignment(container, 2);
+
+		Map<String, Object> senderProps = KafkaTestUtils.producerProps(embeddedKafka);
+		ProducerFactory<Integer, String> pf = new DefaultKafkaProducerFactory<>(senderProps);
+		KafkaTemplate<Integer, String> template = new KafkaTemplate<>(pf);
+		template.setDefaultTopic(topic6);
+		Headers kHeaders = new RecordHeaders();
+		MessageHeaders siHeaders = new MessageHeaders(Collections.singletonMap("foo", "bar"));
+		new DefaultKafkaHeaderMapper().fromHeaders(siHeaders, kHeaders);
+
 		template.sendDefault(1, "{\"bar\":\"baz\"}");
 
-		received = out.receive(10000);
+		Message<?> received = out.receive(10000);
 		assertThat(received).isNotNull();
 
-		headers = received.getHeaders();
+		MessageHeaders headers = received.getHeaders();
 		assertThat(headers.get(KafkaHeaders.RECEIVED_MESSAGE_KEY)).isEqualTo(1);
-		assertThat(headers.get(KafkaHeaders.RECEIVED_TOPIC)).isEqualTo(topic3);
+		assertThat(headers.get(KafkaHeaders.RECEIVED_TOPIC)).isEqualTo(topic6);
 		assertThat(headers.get(KafkaHeaders.RECEIVED_PARTITION_ID)).isEqualTo(0);
-		assertThat(headers.get(KafkaHeaders.OFFSET)).isEqualTo(1L);
+		assertThat(headers.get(KafkaHeaders.OFFSET)).isEqualTo(0L);
 		assertThat((Long) headers.get(KafkaHeaders.RECEIVED_TIMESTAMP)).isGreaterThan(0L);
 		assertThat(headers.get(KafkaHeaders.TIMESTAMP_TYPE)).isEqualTo("CREATE_TIME");