From cbf63d7e90a3d80c0592f571661bc834aba18775 Mon Sep 17 00:00:00 2001
From: Mark Fisher <markfisher@vmware.com>
Date: Wed, 2 Apr 2008 21:34:31 +0000
Subject: [PATCH] MessageEndpointAnnotationPostProcessor now provides the
 'ChannelRegistry' to any Object with a @MessageEndpoint annotation that
 implements ChannelRegistryAware (INT-125).

---
 ...essageEndpointAnnotationPostProcessor.java |  3 ++
 ...eEndpointAnnotationPostProcessorTests.java | 32 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/spring-integration-core/src/main/java/org/springframework/integration/config/MessageEndpointAnnotationPostProcessor.java b/spring-integration-core/src/main/java/org/springframework/integration/config/MessageEndpointAnnotationPostProcessor.java
index b0ea8ce2d0..a8b679732d 100644
--- a/spring-integration-core/src/main/java/org/springframework/integration/config/MessageEndpointAnnotationPostProcessor.java
+++ b/spring-integration-core/src/main/java/org/springframework/integration/config/MessageEndpointAnnotationPostProcessor.java
@@ -104,6 +104,9 @@ public class MessageEndpointAnnotationPostProcessor implements BeanPostProcessor
 		if (endpointAnnotation == null) {
 			return bean;
 		}
+		if (bean instanceof ChannelRegistryAware) {
+			((ChannelRegistryAware) bean).setChannelRegistry(this.messageBus);
+		}
 		MessageHandlerChain handlerChain = this.createHandlerChain(bean);
 		DefaultMessageEndpoint endpoint = new DefaultMessageEndpoint(handlerChain);
 		this.configureInput(bean, beanName, endpointAnnotation, endpoint);
diff --git a/spring-integration-core/src/test/java/org/springframework/integration/endpoint/annotation/MessageEndpointAnnotationPostProcessorTests.java b/spring-integration-core/src/test/java/org/springframework/integration/endpoint/annotation/MessageEndpointAnnotationPostProcessorTests.java
index 5223585256..5af8de4a15 100644
--- a/spring-integration-core/src/test/java/org/springframework/integration/endpoint/annotation/MessageEndpointAnnotationPostProcessorTests.java
+++ b/spring-integration-core/src/test/java/org/springframework/integration/endpoint/annotation/MessageEndpointAnnotationPostProcessorTests.java
@@ -17,6 +17,8 @@
 package org.springframework.integration.endpoint.annotation;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -30,6 +32,8 @@ import org.springframework.integration.annotation.DefaultOutput;
 import org.springframework.integration.annotation.MessageEndpoint;
 import org.springframework.integration.annotation.Polled;
 import org.springframework.integration.bus.MessageBus;
+import org.springframework.integration.channel.ChannelRegistry;
+import org.springframework.integration.channel.ChannelRegistryAware;
 import org.springframework.integration.channel.MessageChannel;
 import org.springframework.integration.channel.SimpleChannel;
 import org.springframework.integration.config.MessageEndpointAnnotationPostProcessor;
@@ -120,6 +124,19 @@ public class MessageEndpointAnnotationPostProcessorTests {
 		new MessageEndpointAnnotationPostProcessor(null);
 	}
 
+	@Test
+	public void testChannelRegistryAwareBean() {
+		MessageBus messageBus = new MessageBus();
+		MessageEndpointAnnotationPostProcessor postProcessor =
+				new MessageEndpointAnnotationPostProcessor(messageBus);
+		ChannelRegistryAwareTestBean testBean = new ChannelRegistryAwareTestBean();
+		assertNull(testBean.getChannelRegistry());
+		postProcessor.postProcessAfterInitialization(testBean, "testBean");
+		ChannelRegistry channelRegistry = testBean.getChannelRegistry();
+		assertNotNull(channelRegistry);
+		assertEquals(messageBus, channelRegistry);
+	}
+
 
 	@MessageEndpoint(defaultOutput="testChannel")
 	private static class PolledAnnotationTestBean {
@@ -160,4 +177,19 @@ public class MessageEndpointAnnotationPostProcessorTests {
 	private static class ConcurrencyAnnotationTestBean {
 	}
 
+
+	@MessageEndpoint(input="inputChannel")
+	private static class ChannelRegistryAwareTestBean implements ChannelRegistryAware {
+
+		private ChannelRegistry channelRegistry;
+
+		public void setChannelRegistry(ChannelRegistry channelRegistry) {
+			this.channelRegistry = channelRegistry;
+		}
+
+		public ChannelRegistry getChannelRegistry() {
+			return this.channelRegistry;
+		}
+	}
+
 }