From 4d3383242e5655631e44dee48ca7060add84f957 Mon Sep 17 00:00:00 2001
From: Phillip Webb <pwebb@vmware.com>
Date: Mon, 3 Jun 2013 12:15:47 -0700
Subject: [PATCH] Support EnvironmentAware ImportSelector/Registrar

Add support for the EnvironmentAware interface with ImportSelector
and ImportBeanDefinitionRegistrar implementations.

Issue: SPR-10602
---
 .../annotation/ConfigurationClassParser.java  |  4 +
 .../ImportBeanDefinitionRegistrar.java        |  3 +-
 .../context/annotation/ImportSelector.java    |  1 +
 .../ImportBeanDefinitionRegistrarTests.java   | 20 +++--
 .../annotation/ImportSelectorTests.java       | 80 ++++++++++++++++---
 5 files changed, 92 insertions(+), 16 deletions(-)

diff --git a/spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassParser.java b/spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassParser.java
index 294d94bb04f..0ac8c31e824 100644
--- a/spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassParser.java
+++ b/spring-context/src/main/java/org/springframework/context/annotation/ConfigurationClassParser.java
@@ -47,6 +47,7 @@ import org.springframework.beans.factory.support.AbstractBeanDefinition;
 import org.springframework.beans.factory.support.BeanDefinitionReader;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanNameGenerator;
+import org.springframework.context.EnvironmentAware;
 import org.springframework.context.ResourceLoaderAware;
 import org.springframework.core.NestedIOException;
 import org.springframework.core.annotation.AnnotationAttributes;
@@ -486,6 +487,9 @@ class ConfigurationClassParser {
 	 */
 	private void invokeAwareMethods(Object importStrategyBean) {
 		if (importStrategyBean instanceof Aware) {
+			if (importStrategyBean instanceof EnvironmentAware) {
+				((EnvironmentAware) importStrategyBean).setEnvironment(this.environment);
+			}
 			if (importStrategyBean instanceof ResourceLoaderAware) {
 				((ResourceLoaderAware) importStrategyBean).setResourceLoader(this.resourceLoader);
 			}
diff --git a/spring-context/src/main/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrar.java b/spring-context/src/main/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrar.java
index 648a8e92113..6795d6172c3 100644
--- a/spring-context/src/main/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrar.java
+++ b/spring-context/src/main/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrar.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2012 the original author or authors.
+ * Copyright 2002-2013 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.
@@ -33,6 +33,7 @@ import org.springframework.core.type.AnnotationMetadata;
  * {@link org.springframework.beans.factory.Aware Aware} interfaces, and their respective
  * methods will be called prior to {@link #registerBeanDefinitions}:
  * <ul>
+ * <li>{@link org.springframework.context.EnvironmentAware}</li>
  * <li>{@link org.springframework.beans.factory.BeanFactoryAware BeanFactoryAware}
  * <li>{@link org.springframework.beans.factory.BeanClassLoaderAware BeanClassLoaderAware}
  * <li>{@link org.springframework.context.ResourceLoaderAware ResourceLoaderAware}
diff --git a/spring-context/src/main/java/org/springframework/context/annotation/ImportSelector.java b/spring-context/src/main/java/org/springframework/context/annotation/ImportSelector.java
index 5d8848001df..b98719f97f1 100644
--- a/spring-context/src/main/java/org/springframework/context/annotation/ImportSelector.java
+++ b/spring-context/src/main/java/org/springframework/context/annotation/ImportSelector.java
@@ -27,6 +27,7 @@ import org.springframework.core.type.AnnotationMetadata;
  * {@link org.springframework.beans.factory.Aware Aware} interfaces, and their respective
  * methods will be called prior to {@link #selectImports}:
  * <ul>
+ * <li>{@link org.springframework.context.EnvironmentAware}</li>
  * <li>{@link org.springframework.beans.factory.BeanFactoryAware BeanFactoryAware}</li>
  * <li>{@link org.springframework.beans.factory.BeanClassLoaderAware BeanClassLoaderAware}</li>
  * <li>{@link org.springframework.context.ResourceLoaderAware ResourceLoaderAware}</li>
diff --git a/spring-context/src/test/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrarTests.java b/spring-context/src/test/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrarTests.java
index 6f55b1f669c..6a45feb0d5b 100644
--- a/spring-context/src/test/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrarTests.java
+++ b/spring-context/src/test/java/org/springframework/context/annotation/ImportBeanDefinitionRegistrarTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2012 the original author or authors.
+ * Copyright 2002-2013 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.
@@ -16,26 +16,27 @@
 
 package org.springframework.context.annotation;
 
-import static org.hamcrest.CoreMatchers.*;
-import static org.junit.Assert.*;
-
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.lang.annotation.Target;
 
 import org.junit.Test;
-
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.BeanClassLoaderAware;
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
+import org.springframework.context.EnvironmentAware;
 import org.springframework.context.MessageSource;
 import org.springframework.context.ResourceLoaderAware;
+import org.springframework.core.env.Environment;
 import org.springframework.core.io.ResourceLoader;
 import org.springframework.core.type.AnnotationMetadata;
 
+import static org.hamcrest.CoreMatchers.*;
+import static org.junit.Assert.*;
+
 /**
  * Integration tests for {@link ImportBeanDefinitionRegistrar}.
  *
@@ -53,6 +54,7 @@ public class ImportBeanDefinitionRegistrarTests {
 		assertThat(SampleRegistrar.beanFactory, is((BeanFactory) context.getBeanFactory()));
 		assertThat(SampleRegistrar.classLoader, is(context.getBeanFactory().getBeanClassLoader()));
 		assertThat(SampleRegistrar.resourceLoader, is(notNullValue()));
+		assertThat(SampleRegistrar.environment, is((Environment) context.getEnvironment()));
 	}
 
 	@Sample
@@ -69,11 +71,12 @@ public class ImportBeanDefinitionRegistrarTests {
 	}
 
 	static class SampleRegistrar implements ImportBeanDefinitionRegistrar, BeanClassLoaderAware, ResourceLoaderAware,
-			BeanFactoryAware {
+			BeanFactoryAware, EnvironmentAware {
 
 		static ClassLoader classLoader;
 		static ResourceLoader resourceLoader;
 		static BeanFactory beanFactory;
+		static Environment environment;
 
 		@Override
 		public void setBeanClassLoader(ClassLoader classLoader) {
@@ -90,6 +93,11 @@ public class ImportBeanDefinitionRegistrarTests {
 			SampleRegistrar.resourceLoader = resourceLoader;
 		}
 
+		@Override
+		public void setEnvironment(Environment environment) {
+			SampleRegistrar.environment = environment;
+		}
+
 		@Override
 		public void registerBeanDefinitions(AnnotationMetadata importingClassMetadata, BeanDefinitionRegistry registry) {
 		}
diff --git a/spring-context/src/test/java/org/springframework/context/annotation/ImportSelectorTests.java b/spring-context/src/test/java/org/springframework/context/annotation/ImportSelectorTests.java
index abdfcdbf6d1..b15e54221c8 100644
--- a/spring-context/src/test/java/org/springframework/context/annotation/ImportSelectorTests.java
+++ b/spring-context/src/test/java/org/springframework/context/annotation/ImportSelectorTests.java
@@ -16,11 +16,6 @@
 
 package org.springframework.context.annotation;
 
-import static org.mockito.Matchers.any;
-import static org.mockito.Matchers.eq;
-import static org.mockito.Mockito.inOrder;
-import static org.mockito.Mockito.spy;
-
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -28,12 +23,27 @@ import java.lang.annotation.Target;
 
 import org.junit.Test;
 import org.mockito.InOrder;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanClassLoaderAware;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.support.DefaultListableBeanFactory;
+import org.springframework.context.EnvironmentAware;
+import org.springframework.context.MessageSource;
+import org.springframework.context.ResourceLoaderAware;
+import org.springframework.context.annotation.ImportBeanDefinitionRegistrarTests.SampleRegistrar;
 import org.springframework.core.Ordered;
 import org.springframework.core.annotation.Order;
+import org.springframework.core.env.Environment;
+import org.springframework.core.io.ResourceLoader;
 import org.springframework.core.type.AnnotationMetadata;
 
+import static org.hamcrest.CoreMatchers.*;
+import static org.junit.Assert.*;
+import static org.mockito.Matchers.*;
+import static org.mockito.Mockito.*;
+
 /**
  * Tests for {@link ImportSelector} and {@link DeferredImportSelector}.
  *
@@ -50,10 +60,62 @@ public class ImportSelectorTests {
 		context.refresh();
 		context.getBean(Config.class);
 		InOrder ordered = inOrder(beanFactory);
-		ordered.verify(beanFactory).registerBeanDefinition(eq("a"), any(BeanDefinition.class));
-		ordered.verify(beanFactory).registerBeanDefinition(eq("b"), any(BeanDefinition.class));
-		ordered.verify(beanFactory).registerBeanDefinition(eq("d"), any(BeanDefinition.class));
-		ordered.verify(beanFactory).registerBeanDefinition(eq("c"), any(BeanDefinition.class));
+		ordered.verify(beanFactory).registerBeanDefinition(eq("a"), (BeanDefinition) anyObject());
+		ordered.verify(beanFactory).registerBeanDefinition(eq("b"), (BeanDefinition) anyObject());
+		ordered.verify(beanFactory).registerBeanDefinition(eq("d"), (BeanDefinition) anyObject());
+		ordered.verify(beanFactory).registerBeanDefinition(eq("c"), (BeanDefinition) anyObject());
+	}
+
+	@Test
+	public void invokeAwareMethodsInImportSelector() {
+
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(AwareConfig.class);
+		context.getBean(MessageSource.class);
+
+		assertThat(SampleRegistrar.beanFactory, is((BeanFactory) context.getBeanFactory()));
+		assertThat(SampleRegistrar.classLoader, is(context.getBeanFactory().getBeanClassLoader()));
+		assertThat(SampleRegistrar.resourceLoader, is(notNullValue()));
+		assertThat(SampleRegistrar.environment, is((Environment) context.getEnvironment()));
+	}
+
+	@Configuration
+	@Import(SampleImportSelector.class)
+	static class AwareConfig {
+
+	}
+
+	static class SampleImportSelector implements ImportSelector, BeanClassLoaderAware, ResourceLoaderAware,
+			BeanFactoryAware, EnvironmentAware {
+
+		static ClassLoader classLoader;
+		static ResourceLoader resourceLoader;
+		static BeanFactory beanFactory;
+		static Environment environment;
+
+		@Override
+		public void setBeanClassLoader(ClassLoader classLoader) {
+			SampleRegistrar.classLoader = classLoader;
+		}
+
+		@Override
+		public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
+			SampleRegistrar.beanFactory = beanFactory;
+		}
+
+		@Override
+		public void setResourceLoader(ResourceLoader resourceLoader) {
+			SampleRegistrar.resourceLoader = resourceLoader;
+		}
+
+		@Override
+		public void setEnvironment(Environment environment) {
+			SampleRegistrar.environment = environment;
+		}
+
+		@Override
+		public String[] selectImports(AnnotationMetadata importingClassMetadata) {
+			return new String[] {};
+		}
 	}
 
 	@Sample