Support getting environment from other events.

Author: youngm

spring-boot/src/main/java/org/springframework/boot/system/ApplicationPidFileWriter.java

@@ -18,6 +18,8 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -34,6 +36,7 @@
 import org.springframework.context.ApplicationListener;
 import org.springframework.core.Ordered;
 import org.springframework.core.env.Environment;
+import org.springframework.core.env.EnvironmentCapable;
 import org.springframework.util.Assert;
 
 /**
@@ -49,7 +52,8 @@
  * <p>
  * Note: access to the Spring {@link Environment} is only possible when the
  * {@link #setTriggerEventType(Class) triggerEventType} is set to
- * {@link ApplicationEnvironmentPreparedEvent} or {@link ApplicationPreparedEvent}.
+ * {@link ApplicationEnvironmentPreparedEvent} or any event that has a getApplicationContext
+ * method.
  *
  * @author Jakub Kubrynski
  * @author Dave Syer
@@ -227,11 +231,19 @@ private Environment getEnvironment(SpringApplicationEvent event) {
 			if (event instanceof ApplicationEnvironmentPreparedEvent) {
 				return ((ApplicationEnvironmentPreparedEvent) event).getEnvironment();
 			}
-			if (event instanceof ApplicationPreparedEvent) {
-				return ((ApplicationPreparedEvent) event).getApplicationContext()
-						.getEnvironment();
+
+			try {
+				Method getApplicationContextMethod = event.getClass().getMethod("getApplicationContext", (Class[]) null);
+				return ((EnvironmentCapable) getApplicationContextMethod.invoke(event, (Object[]) null)).getEnvironment();
+			}
+			catch (NoSuchMethodException | IllegalAccessException e) {
+				logger.warn("Event '" + event.getClass() + "' does not provide a 'getApplicationContext()' method.  Ignoring Environment config.");
+				return null;
+			}
+			catch (InvocationTargetException e) {
+				logger.warn("Unexpected error getting application context from event.  Ignoring Environment config.", e.getTargetException());
+				return null;
 			}
-			return null;
 		}
 
 	}

spring-boot/src/test/java/org/springframework/boot/system/ApplicationPidFileWriterTests.java

@@ -29,6 +29,7 @@
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.context.event.ApplicationEnvironmentPreparedEvent;
 import org.springframework.boot.context.event.ApplicationPreparedEvent;
+import org.springframework.boot.context.event.ApplicationReadyEvent;
 import org.springframework.boot.context.event.ApplicationStartedEvent;
 import org.springframework.boot.context.event.SpringApplicationEvent;
 import org.springframework.context.ConfigurableApplicationContext;
@@ -99,7 +100,7 @@ public void overridePidFileWithSpring() throws Exception {
 	}
 
 	@Test
-	public void differentEventTypes() throws Exception {
+	public void tryEnvironmentPreparedEvent() throws Exception {
 		File file = this.temporaryFolder.newFile();
 		SpringApplicationEvent event = createEnvironmentPreparedEvent("spring.pid.file",
 				file.getAbsolutePath());
@@ -111,6 +112,19 @@ public void differentEventTypes() throws Exception {
 		assertThat(FileCopyUtils.copyToString(new FileReader(file))).isNotEmpty();
 	}
 
+	@Test
+	public void tryReadyEvent() throws Exception {
+		File file = this.temporaryFolder.newFile();
+		SpringApplicationEvent event = createReadyEvent("spring.pid.file",
+				file.getAbsolutePath());
+		ApplicationPidFileWriter listener = new ApplicationPidFileWriter();
+		listener.onApplicationEvent(event);
+		assertThat(FileCopyUtils.copyToString(new FileReader(file))).isEmpty();
+		listener.setTriggerEventType(ApplicationReadyEvent.class);
+		listener.onApplicationEvent(event);
+		assertThat(FileCopyUtils.copyToString(new FileReader(file))).isNotEmpty();
+	}
+
 	@Test
 	public void withNoEnvironment() throws Exception {
 		File file = this.temporaryFolder.newFile();
@@ -170,6 +184,16 @@ private SpringApplicationEvent createPreparedEvent(String propName,
 				context);
 	}
 
+	private SpringApplicationEvent createReadyEvent(String propName,
+			String propValue) {
+		ConfigurableEnvironment environment = createEnvironment(propName, propValue);
+		ConfigurableApplicationContext context = mock(
+				ConfigurableApplicationContext.class);
+		given(context.getEnvironment()).willReturn(environment);
+		return new ApplicationReadyEvent(new SpringApplication(), new String[] {},
+				context);
+	}
+
 	private ConfigurableEnvironment createEnvironment(String propName, String propValue) {
 		MockPropertySource propertySource = mockPropertySource(propName, propValue);
 		ConfigurableEnvironment environment = new StandardEnvironment();