From d4251125fc829d1e7640a52c8d9d821854353455 Mon Sep 17 00:00:00 2001
From: Stephane Nicoll <snicoll@pivotal.io>
Date: Wed, 12 Apr 2017 14:30:42 +0200
Subject: [PATCH] Make sure Admin MBean works also with webflux

This commit improves `SpringApplicationAdminMXBeanRegistrar` so that it
work with traditional Servlet-based app but also with Spring WebFlux.

Closes gh-8533
---
 ...SpringApplicationAdminMXBeanRegistrar.java | 22 +++++++++++++------
 ...gApplicationAdminMXBeanRegistrarTests.java |  4 ++--
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/spring-boot/src/main/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrar.java b/spring-boot/src/main/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrar.java
index b0972c2ef5..b725654511 100644
--- a/spring-boot/src/main/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrar.java
+++ b/spring-boot/src/main/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrar.java
@@ -29,12 +29,12 @@ import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.boot.context.event.ApplicationReadyEvent;
-import org.springframework.boot.web.servlet.context.ServletWebServerApplicationContext;
+import org.springframework.boot.web.context.WebServerInitializedEvent;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
-import org.springframework.context.ApplicationListener;
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.EnvironmentAware;
+import org.springframework.context.event.EventListener;
 import org.springframework.core.env.Environment;
 import org.springframework.core.env.StandardEnvironment;
 import org.springframework.util.Assert;
@@ -49,7 +49,7 @@ import org.springframework.util.Assert;
  */
 public class SpringApplicationAdminMXBeanRegistrar
 		implements ApplicationContextAware, EnvironmentAware, InitializingBean,
-		DisposableBean, ApplicationListener<ApplicationReadyEvent> {
+		DisposableBean {
 
 	private static final Log logger = LogFactory.getLog(SpringApplicationAdmin.class);
 
@@ -61,6 +61,8 @@ public class SpringApplicationAdminMXBeanRegistrar
 
 	private boolean ready = false;
 
+	private boolean embeddedWebApplication = false;
+
 	public SpringApplicationAdminMXBeanRegistrar(String name)
 			throws MalformedObjectNameException {
 		this.objectName = new ObjectName(name);
@@ -79,13 +81,20 @@ public class SpringApplicationAdminMXBeanRegistrar
 		this.environment = environment;
 	}
 
-	@Override
-	public void onApplicationEvent(ApplicationReadyEvent event) {
+	@EventListener
+	public void onApplicationReadyEvent(ApplicationReadyEvent event) {
 		if (this.applicationContext.equals(event.getApplicationContext())) {
 			this.ready = true;
 		}
 	}
 
+	@EventListener
+	public void onWebServerInitializedEvent(WebServerInitializedEvent event) {
+		if (this.applicationContext.equals(event.getApplicationContext())) {
+			this.embeddedWebApplication = true;
+		}
+	}
+
 	@Override
 	public void afterPropertiesSet() throws Exception {
 		MBeanServer server = ManagementFactory.getPlatformMBeanServer();
@@ -110,8 +119,7 @@ public class SpringApplicationAdminMXBeanRegistrar
 
 		@Override
 		public boolean isEmbeddedWebApplication() {
-			return (SpringApplicationAdminMXBeanRegistrar.this.applicationContext != null
-					&& SpringApplicationAdminMXBeanRegistrar.this.applicationContext instanceof ServletWebServerApplicationContext);
+			return SpringApplicationAdminMXBeanRegistrar.this.embeddedWebApplication;
 		}
 
 		@Override
diff --git a/spring-boot/src/test/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrarTests.java b/spring-boot/src/test/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrarTests.java
index 817e507b57..fbe645b2c0 100644
--- a/spring-boot/src/test/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrarTests.java
+++ b/spring-boot/src/test/java/org/springframework/boot/admin/SpringApplicationAdminMXBeanRegistrarTests.java
@@ -99,10 +99,10 @@ public class SpringApplicationAdminMXBeanRegistrarTests {
 		ConfigurableApplicationContext context = mock(
 				ConfigurableApplicationContext.class);
 		registrar.setApplicationContext(context);
-		registrar.onApplicationEvent(new ApplicationReadyEvent(new SpringApplication(),
+		registrar.onApplicationReadyEvent(new ApplicationReadyEvent(new SpringApplication(),
 				null, mock(ConfigurableApplicationContext.class)));
 		assertThat(isApplicationReady(registrar)).isFalse();
-		registrar.onApplicationEvent(
+		registrar.onApplicationReadyEvent(
 				new ApplicationReadyEvent(new SpringApplication(), null, context));
 		assertThat(isApplicationReady(registrar)).isTrue();
 	}