From e6afc490b088ff537a64753bc6cafbfd2748e0c7 Mon Sep 17 00:00:00 2001
From: Stephane Nicoll <snicoll@vmware.com>
Date: Thu, 4 May 2023 16:56:11 +0200
Subject: [PATCH] Expose ApplicationContext when using SpringApplication#from

This commit returns the running application context when running an
application via SpringApplication#from so that it is consistent with
the regular SpringApplication#run.

Closes gh-35203
---
 .../boot/SpringApplication.java               | 31 ++++++++++++++++---
 .../boot/SpringApplicationTests.java          |  8 +++++
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
index c420f86e54..28ea16bf88 100644
--- a/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
+++ b/spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java
@@ -1453,14 +1453,35 @@ public class SpringApplication {
 		/**
 		 * Run the application using the given args.
 		 * @param args the main method args
+		 * @return the running {@link ApplicationContext}
 		 */
-		public void run(String... args) {
-			withHook(this::getRunListener, () -> this.main.accept(args));
+		public ConfigurableApplicationContext run(String... args) {
+			ContextLoaderHook hook = new ContextLoaderHook(this.sources);
+			withHook(hook, () -> this.main.accept(args));
+			return hook.applicationContext;
 		}
 
-		private SpringApplicationRunListener getRunListener(SpringApplication springApplication) {
-			springApplication.addPrimarySources(this.sources);
-			return null;
+		private static class ContextLoaderHook implements SpringApplicationHook {
+
+			private final Set<Class<?>> sources;
+
+			private ConfigurableApplicationContext applicationContext;
+
+			ContextLoaderHook(Set<Class<?>> sources) {
+				this.sources = sources;
+			}
+
+			@Override
+			public SpringApplicationRunListener getRunListener(SpringApplication springApplication) {
+				springApplication.addPrimarySources(this.sources);
+				return new SpringApplicationRunListener() {
+					@Override
+					public void contextPrepared(ConfigurableApplicationContext context) {
+						ContextLoaderHook.this.applicationContext = context;
+					}
+				};
+			}
+
 		}
 
 	}
diff --git a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java
index 5d0ef06b5c..22d3ffd6fa 100644
--- a/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java
+++ b/spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java
@@ -1368,6 +1368,14 @@ class SpringApplicationTests {
 		ExampleAdditionalConfig.local.set(null);
 	}
 
+	@Test
+	void fromReturnsApplicationContext() {
+		ConfigurableApplicationContext context = SpringApplication.from(ExampleFromMainMethod::main)
+			.with(ExampleAdditionalConfig.class)
+			.run();
+		assertThat(context).isNotNull();
+	}
+
 	private <S extends AvailabilityState> ArgumentMatcher<ApplicationEvent> isAvailabilityChangeEventWithState(
 			S state) {
 		return (argument) -> (argument instanceof AvailabilityChangeEvent<?>)