From 6eff068ca1f3047dda9cd78f93d29dce9824a65c Mon Sep 17 00:00:00 2001
From: Oleg Zhurakousky <ozhurakousky@vmware.com>
Date: Mon, 16 May 2011 22:08:16 -0400
Subject: [PATCH] INT-1903 polishing

---
 .../config/IdGeneratorConfigurer.java         | 22 +++++++++----------
 .../core/MessageIdGenerationTests.java        |  3 +--
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/spring-integration-core/src/main/java/org/springframework/integration/config/IdGeneratorConfigurer.java b/spring-integration-core/src/main/java/org/springframework/integration/config/IdGeneratorConfigurer.java
index 22e6c099a0..7da7505c54 100644
--- a/spring-integration-core/src/main/java/org/springframework/integration/config/IdGeneratorConfigurer.java
+++ b/spring-integration-core/src/main/java/org/springframework/integration/config/IdGeneratorConfigurer.java
@@ -21,7 +21,6 @@ import java.lang.reflect.Field;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
-import org.springframework.beans.factory.BeanDefinitionStoreException;
 import org.springframework.beans.factory.NoSuchBeanDefinitionException;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationListener;
@@ -30,6 +29,7 @@ import org.springframework.context.event.ContextClosedEvent;
 import org.springframework.context.event.ContextRefreshedEvent;
 import org.springframework.integration.MessageHeaders;
 import org.springframework.integration.MessageHeaders.IdGenerator;
+import org.springframework.util.Assert;
 import org.springframework.util.ReflectionUtils;
 import org.springframework.util.StringUtils;
 
@@ -45,21 +45,19 @@ public final class IdGeneratorConfigurer implements ApplicationListener<Applicat
 
 
 	public void onApplicationEvent(ApplicationContextEvent event) {
-		ApplicationContext contex = event.getApplicationContext();
+		ApplicationContext context = event.getApplicationContext();
 		if (event instanceof ContextRefreshedEvent){
-			if (!StringUtils.hasText(IdGeneratorConfigurer.generatorContextId)){
-				
-				if (this.setIdGenerator(contex)){
-					IdGeneratorConfigurer.generatorContextId = contex.getId();
-				}	
-			}
-			else if (contex.getBeanNamesForType(IdGenerator.class).length > 0) {
-				throw new BeanDefinitionStoreException(
-					"'MessageHeaders.idGenerator' has already been set and can not be set again");
+			boolean contextHasIdGenerator = context.getBeanNamesForType(IdGenerator.class).length > 0;
+			if (contextHasIdGenerator) {
+				Assert.state(!StringUtils.hasText(IdGeneratorConfigurer.generatorContextId),
+						"'MessageHeaders.idGenerator' has already been set and can not be set again");
+				if (this.setIdGenerator(context)) {
+					IdGeneratorConfigurer.generatorContextId = context.getId();
+				}
 			}
 		}
 		else if (event instanceof ContextClosedEvent){
-			if (contex.getId().equals(IdGeneratorConfigurer.generatorContextId)){
+			if (context.getId().equals(IdGeneratorConfigurer.generatorContextId)){
 				this.unsetIdGenerator();
 				IdGeneratorConfigurer.generatorContextId = null;
 			}
diff --git a/spring-integration-core/src/test/java/org/springframework/integration/core/MessageIdGenerationTests.java b/spring-integration-core/src/test/java/org/springframework/integration/core/MessageIdGenerationTests.java
index 8051b80806..1e9506f5b5 100644
--- a/spring-integration-core/src/test/java/org/springframework/integration/core/MessageIdGenerationTests.java
+++ b/spring-integration-core/src/test/java/org/springframework/integration/core/MessageIdGenerationTests.java
@@ -25,7 +25,6 @@ import java.util.UUID;
 import org.junit.Ignore;
 import org.junit.Test;
 
-import org.springframework.beans.factory.BeanDefinitionStoreException;
 import org.springframework.context.support.ClassPathXmlApplicationContext;
 import org.springframework.context.support.GenericXmlApplicationContext;
 import org.springframework.integration.MessageChannel;
@@ -125,7 +124,7 @@ public class MessageIdGenerationTests {
 		parent.close();
 	}
 	
-	@Test(expected=BeanDefinitionStoreException.class)
+	@Test(expected=IllegalStateException.class)
 	public void testCustomIdGenerationWithParentChileIndependentCreationTwoChildrenTwoRegistrars(){
 		ClassPathXmlApplicationContext parent = new ClassPathXmlApplicationContext("MessageIdGenerationTests-context-withGenerator.xml", this.getClass());