From e2262284f1065a664fb52a3fa24a7fce9f01f244 Mon Sep 17 00:00:00 2001
From: Phillip Webb <pwebb@vmware.com>
Date: Mon, 10 Apr 2023 20:15:57 -0700
Subject: [PATCH] Add SpringApplication.from(...) support

Add a `SpringApplication.from(...)` method which can be used
to augment an existing application with addition `@Configuration`.

Closes gh-35019
---
 .../boot/SpringApplication.java               | 61 +++++++++++++++++++
 .../boot/SpringApplicationTests.java          | 29 +++++++++
 2 files changed, 90 insertions(+)

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 2ebe1bd5a5..c420f86e54 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
@@ -88,6 +88,7 @@ import org.springframework.util.CollectionUtils;
 import org.springframework.util.ObjectUtils;
 import org.springframework.util.ReflectionUtils;
 import org.springframework.util.StringUtils;
+import org.springframework.util.function.ThrowingConsumer;
 import org.springframework.util.function.ThrowingSupplier;
 
 /**
@@ -1357,6 +1358,22 @@ public class SpringApplication {
 		return exitCode;
 	}
 
+	/**
+	 * Create an application from an existing {@code main} method that can run with
+	 * additional {@code @Configuration} or bean classes. This method can be helpful when
+	 * writing a test harness that needs to start an application with additional
+	 * configuration.
+	 * @param main the main method entry point that runs the {@link SpringApplication}
+	 * @return an {@link SpringApplication.Augmented} instance that can be used to add
+	 * configuration and run the application.
+	 * @since 3.1.0
+	 * @see #withHook(SpringApplicationHook, Runnable)
+	 */
+	public static SpringApplication.Augmented from(ThrowingConsumer<String[]> main) {
+		Assert.notNull(main, "Main must not be null");
+		return new Augmented(main, Collections.emptySet());
+	}
+
 	/**
 	 * Perform the given action with the given {@link SpringApplicationHook} attached if
 	 * the action triggers an {@link SpringApplication#run(String...) application run}.
@@ -1404,6 +1421,50 @@ public class SpringApplication {
 		return new LinkedHashSet<>(list);
 	}
 
+	/**
+	 * Used to configure and run an augmented {@link SpringApplication} where additional
+	 * configuration should be applied.
+	 *
+	 * @since 3.1.0
+	 */
+	public static class Augmented {
+
+		private final ThrowingConsumer<String[]> main;
+
+		private final Set<Class<?>> sources;
+
+		Augmented(ThrowingConsumer<String[]> main, Set<Class<?>> sources) {
+			this.main = main;
+			this.sources = Set.copyOf(sources);
+		}
+
+		/**
+		 * Return a new {@link SpringApplication.Augmented} instance with additional
+		 * sources that should be applied when the application runs.
+		 * @param sources the sources that should be applied
+		 * @return a new {@link SpringApplication.Augmented} instance
+		 */
+		public Augmented with(Class<?>... sources) {
+			LinkedHashSet<Class<?>> merged = new LinkedHashSet<>(this.sources);
+			merged.addAll(Arrays.asList(sources));
+			return new Augmented(this.main, merged);
+		}
+
+		/**
+		 * Run the application using the given args.
+		 * @param args the main method args
+		 */
+		public void run(String... args) {
+			withHook(this::getRunListener, () -> this.main.accept(args));
+		}
+
+		private SpringApplicationRunListener getRunListener(SpringApplication springApplication) {
+			springApplication.addPrimarySources(this.sources);
+			return null;
+		}
+
+	}
+
 	/**
 	 * {@link BeanFactoryPostProcessor} to re-order our property sources below any
 	 * {@code @PropertySource} items added by the {@link ConfigurationClassPostProcessor}.
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 696f287bcd..5d0ef06b5c 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
@@ -1360,6 +1360,14 @@ class SpringApplicationTests {
 		}
 	}
 
+	@Test
+	void fromRunsWithAdditionalSources() {
+		assertThat(ExampleAdditionalConfig.local.get()).isNull();
+		SpringApplication.from(ExampleFromMainMethod::main).with(ExampleAdditionalConfig.class).run();
+		assertThat(ExampleAdditionalConfig.local.get()).isNotNull();
+		ExampleAdditionalConfig.local.set(null);
+	}
+
 	private <S extends AvailabilityState> ArgumentMatcher<ApplicationEvent> isAvailabilityChangeEventWithState(
 			S state) {
 		return (argument) -> (argument instanceof AvailabilityChangeEvent<?>)
@@ -1922,4 +1930,25 @@ class SpringApplicationTests {
 
 	}
 
+	static class ExampleFromMainMethod {
+
+		static void main(String[] args) {
+			SpringApplication application = new SpringApplication(ExampleConfig.class);
+			application.setWebApplicationType(WebApplicationType.NONE);
+			application.run(args);
+		}
+
+	}
+
+	@Configuration
+	static class ExampleAdditionalConfig {
+
+		static ThreadLocal<ExampleAdditionalConfig> local = new ThreadLocal<>();
+
+		ExampleAdditionalConfig() {
+			local.set(this);
+		}
+
+	}
+
 }