Allow the configuration of active profiles in SpringApplication.Augmented

Closes gh-36660

Author: mhalbritter

spring-boot-project/spring-boot/src/main/java/org/springframework/boot/SpringApplication.java

@@ -1430,7 +1430,7 @@ public static int exit(ApplicationContext context, ExitCodeGenerator... exitCode
 	 */
 	public static SpringApplication.Augmented from(ThrowingConsumer<String[]> main) {
 		Assert.notNull(main, "Main must not be null");
-		return new Augmented(main, Collections.emptySet());
+		return new Augmented(main, Collections.emptySet(), Collections.emptySet());
 	}
 
 	/**
@@ -1492,9 +1492,12 @@ public static class Augmented {
 
 		private final Set<Class<?>> sources;
 
-		Augmented(ThrowingConsumer<String[]> main, Set<Class<?>> sources) {
+		private final Set<String> additionalProfiles;
+
+		Augmented(ThrowingConsumer<String[]> main, Set<Class<?>> sources, Set<String> additionalProfiles) {
 			this.main = main;
 			this.sources = Set.copyOf(sources);
+			this.additionalProfiles = additionalProfiles;
 		}
 
 		/**
@@ -1506,7 +1509,20 @@ public static class Augmented {
 		public Augmented with(Class<?>... sources) {
 			LinkedHashSet<Class<?>> merged = new LinkedHashSet<>(this.sources);
 			merged.addAll(Arrays.asList(sources));
-			return new Augmented(this.main, merged);
+			return new Augmented(this.main, merged, this.additionalProfiles);
+		}
+
+		/**
+		 * Return a new {@link SpringApplication.Augmented} instance with additional
+		 * profiles that should be applied when the application runs.
+		 * @param profiles the profiles that should be applied
+		 * @return a new {@link SpringApplication.Augmented} instance
+		 * @since 3.4.0
+		 */
+		public Augmented withAdditionalProfiles(String... profiles) {
+			Set<String> merged = new LinkedHashSet<>(this.additionalProfiles);
+			merged.addAll(Arrays.asList(profiles));
+			return new Augmented(this.main, this.sources, merged);
 		}
 
 		/**
@@ -1518,6 +1534,7 @@ public SpringApplication.Running run(String... args) {
 			RunListener runListener = new RunListener();
 			SpringApplicationHook hook = new SingleUseSpringApplicationHook((springApplication) -> {
 				springApplication.addPrimarySources(this.sources);
+				springApplication.setAdditionalProfiles(this.additionalProfiles.toArray(String[]::new));
 				return runListener;
 			});
 			withHook(hook, () -> this.main.accept(args));

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/SpringApplicationTests.java

@@ -92,6 +92,7 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Lazy;
+import org.springframework.context.annotation.Profile;
 import org.springframework.context.event.ApplicationEventMulticaster;
 import org.springframework.context.event.ContextRefreshedEvent;
 import org.springframework.context.event.SimpleApplicationEventMulticaster;
@@ -1409,7 +1410,8 @@ public void contextLoaded(ConfigurableApplicationContext context) {
 		then(listener).should(never()).onApplicationEvent(any(ApplicationFailedEvent.class));
 	}
 
-	@Test // gh-32555
+	@Test
+	// gh-32555
 	void shouldUseAotInitializer() {
 		SpringApplication application = new SpringApplication(ExampleAotProcessedMainClass.class);
 		application.setWebApplicationType(WebApplicationType.NONE);
@@ -1468,6 +1470,17 @@ void fromWithMultipleApplicationsOnlyAppliesAdditionalSourcesOnce() {
 		assertThatNoException().isThrownBy(() -> this.context.getBean(SingleUseAdditionalConfig.class));
 	}
 
+	@Test
+	void fromAppliesProfiles() {
+		this.context = SpringApplication.from(ExampleFromMainMethod::main)
+			.with(ProfileConfig.class)
+			.withAdditionalProfiles("custom")
+			.run()
+			.getApplicationContext();
+		assertThat(this.context).isNotNull();
+		assertThat(this.context.getBeanProvider(Example.class).getIfAvailable()).isNotNull();
+	}
+
 	@Test
 	void shouldStartDaemonThreadIfKeepAliveIsEnabled() {
 		SpringApplication application = new SpringApplication(ExampleConfig.class);
@@ -2159,4 +2172,15 @@ static class SingleUseAdditionalConfig {
 
 	}
 
+	@Configuration
+	static class ProfileConfig {
+
+		@Bean
+		@Profile("custom")
+		Example example() {
+			return new Example();
+		}
+
+	}
+
 }