From 63e715a685ac8f53d80629a6728b168b1faab7c0 Mon Sep 17 00:00:00 2001
From: Artem Bilan <abilan@pivotal.io>
Date: Mon, 30 Jan 2017 13:19:01 -0500
Subject: [PATCH] DSL-140: Improve Logic for outputChannel in DSL

Fixes spring-projects/spring-integration-java-dsl#140

Allow to provide any custom `MessageProducer` impl for the `.handle()`, not only `AbstractReplyProducingMessageHandler` extension

Polishing
---
 .../dsl/IntegrationFlowDefinition.java        | 27 +++++++----
 .../dsl/manualflow/ManualFlowTests.java       | 46 ++++++++++++++++++-
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/spring-integration-core/src/main/java/org/springframework/integration/dsl/IntegrationFlowDefinition.java b/spring-integration-core/src/main/java/org/springframework/integration/dsl/IntegrationFlowDefinition.java
index c8a2e5239a..a11cb21c18 100644
--- a/spring-integration-core/src/main/java/org/springframework/integration/dsl/IntegrationFlowDefinition.java
+++ b/spring-integration-core/src/main/java/org/springframework/integration/dsl/IntegrationFlowDefinition.java
@@ -45,6 +45,7 @@ import org.springframework.integration.config.ConsumerEndpointFactoryBean;
 import org.springframework.integration.config.SourcePollingChannelAdapterFactoryBean;
 import org.springframework.integration.context.IntegrationContextUtils;
 import org.springframework.integration.core.GenericSelector;
+import org.springframework.integration.core.MessageProducer;
 import org.springframework.integration.core.MessageSelector;
 import org.springframework.integration.dsl.channel.MessageChannelSpec;
 import org.springframework.integration.dsl.channel.WireTapSpec;
@@ -117,7 +118,7 @@ public abstract class IntegrationFlowDefinition<B extends IntegrationFlowDefinit
 
 	private static final SpelExpressionParser PARSER = new SpelExpressionParser();
 
-	private static final Set<MessageHandler> REFERENCED_REPLY_PRODUCERS = new HashSet<>();
+	private static final Set<MessageProducer> REFERENCED_REPLY_PRODUCERS = new HashSet<>();
 
 	protected final Set<Object> integrationComponents = new LinkedHashSet<>();
 
@@ -2821,12 +2822,22 @@ public abstract class IntegrationFlowDefinition<B extends IntegrationFlowDefinit
 					currentComponent = extractProxyTarget(currentComponent);
 				}
 
-				if (currentComponent instanceof AbstractMessageProducingHandler) {
-					AbstractMessageProducingHandler messageProducer =
-							(AbstractMessageProducingHandler) currentComponent;
+				if (currentComponent instanceof MessageProducer) {
+					MessageProducer messageProducer =
+							(MessageProducer) currentComponent;
 					checkReuse(messageProducer);
 					if (channelName != null) {
-						messageProducer.setOutputChannelName(channelName);
+						if (messageProducer instanceof AbstractMessageProducingHandler) {
+							((AbstractMessageProducingHandler) messageProducer).setOutputChannelName(channelName);
+						}
+						else {
+							throw new BeanCreationException("The 'currentComponent' (" + currentComponent
+									+ ") must extend 'AbstractMessageProducingHandler' "
+									+ "for message channel resolution by name.\n"
+									+ "Your handler should extend 'AbstractMessageProducingHandler', "
+									+ "its subclass 'AbstractReplyProducingMessageHandler', or you should "
+									+ "reference a 'MessageChannel' bean instead of its name.");
+						}
 					}
 					else {
 						messageProducer.setOutputChannel(outputChannel);
@@ -2927,10 +2938,10 @@ public abstract class IntegrationFlowDefinition<B extends IntegrationFlowDefinit
 		}
 	}
 
-	private void checkReuse(AbstractMessageProducingHandler replyHandler) {
+	private void checkReuse(MessageProducer replyHandler) {
 		Assert.isTrue(!REFERENCED_REPLY_PRODUCERS.contains(replyHandler),
-				"An AbstractMessageProducingHandler may only be referenced once ("
-						+ replyHandler.getComponentName()
+				"A reply MessageProducer may only be referenced once ("
+						+ replyHandler
 						+ ") - use @Scope(ConfigurableBeanFactory.SCOPE_PROTOTYPE) on @Bean definition.");
 		REFERENCED_REPLY_PRODUCERS.add(replyHandler);
 	}
diff --git a/spring-integration-core/src/test/java/org/springframework/integration/dsl/manualflow/ManualFlowTests.java b/spring-integration-core/src/test/java/org/springframework/integration/dsl/manualflow/ManualFlowTests.java
index dc51790efe..880bef710e 100644
--- a/spring-integration-core/src/test/java/org/springframework/integration/dsl/manualflow/ManualFlowTests.java
+++ b/spring-integration-core/src/test/java/org/springframework/integration/dsl/manualflow/ManualFlowTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2016 the original author or authors.
+ * Copyright 2016-2017 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.
@@ -42,6 +42,7 @@ import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Scope;
 import org.springframework.integration.channel.QueueChannel;
 import org.springframework.integration.config.EnableIntegration;
+import org.springframework.integration.core.MessageProducer;
 import org.springframework.integration.core.MessagingTemplate;
 import org.springframework.integration.dsl.IntegrationFlow;
 import org.springframework.integration.dsl.IntegrationFlowAdapter;
@@ -51,6 +52,9 @@ import org.springframework.integration.dsl.context.IntegrationFlowContext;
 import org.springframework.integration.dsl.context.IntegrationFlowRegistration;
 import org.springframework.integration.handler.AbstractReplyProducingMessageHandler;
 import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.MessagingException;
 import org.springframework.messaging.PollableChannel;
 import org.springframework.messaging.support.GenericMessage;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
@@ -205,6 +209,46 @@ public class ManualFlowTests {
 		}
 	}
 
+	@Test
+	public void testMessageProducerForOutputChannel() {
+		class MessageProducingHandler implements MessageHandler, MessageProducer {
+
+			private MessageChannel outputChannel;
+
+			@Override
+			public void setOutputChannel(MessageChannel outputChannel) {
+				this.outputChannel = outputChannel;
+			}
+
+			@Override
+			public MessageChannel getOutputChannel() {
+				return this.outputChannel;
+			}
+
+			@Override
+			public void handleMessage(Message<?> message) throws MessagingException {
+				this.outputChannel.send(message);
+			}
+
+		}
+
+
+		PollableChannel resultChannel = new QueueChannel();
+
+		IntegrationFlowRegistration flowRegistration =
+				this.integrationFlowContext.registration(flow ->
+						flow.handle(new MessageProducingHandler())
+								.channel(resultChannel))
+						.register();
+
+		this.integrationFlowContext.messagingTemplateFor(flowRegistration.getId())
+				.send(new GenericMessage<>("test"));
+
+		Message<?> receive = resultChannel.receive(1000);
+		assertNotNull(receive);
+		assertEquals("test", receive.getPayload());
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class RootConfiguration {