Support "args" param in "@SpringBootTest"

spring-aot-test/src/main/java/org/springframework/aot/test/boot/SpringBootAotTestContextProcessor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,15 +16,19 @@
 
 package org.springframework.aot.test.boot;
 
+import java.util.Arrays;
+
 import com.squareup.javapoet.ClassName;
 import com.squareup.javapoet.CodeBlock;
 import com.squareup.javapoet.CodeBlock.Builder;
 
+import org.springframework.aot.context.bootstrap.generator.bean.support.MultiCodeBlock;
 import org.springframework.aot.test.context.bootstrap.generator.AotTestContextProcessor;
 import org.springframework.boot.WebApplicationType;
 import org.springframework.boot.test.context.ReactiveWebMergedContextConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.boot.test.context.SpringBootTest.WebEnvironment;
+import org.springframework.boot.test.context.SpringBootTestArgsAccessor;
 import org.springframework.boot.test.context.SpringBootTestContextBootstrapper;
 import org.springframework.context.support.GenericApplicationContext;
 import org.springframework.core.annotation.MergedAnnotations;
@@ -62,16 +66,22 @@ public GenericApplicationContext prepareTestContext(MergedContextConfiguration c
 
 	@Override
 	public CodeBlock writeInstanceSupplier(MergedContextConfiguration config, ClassName applicationContextInitializer) {
+		String[] args = SpringBootTestArgsAccessor.get(config.getContextCustomizers());
+
 		Builder code = CodeBlock.builder();
 		code.add("() -> new $T($T.class", AotSpringBootConfigContextLoader.class, applicationContextInitializer);
 		WebApplicationType webApplicationType = detectWebApplicationType(config);
-		if (webApplicationType.equals(WebApplicationType.NONE)) {
-			code.add(")");
-		}
-		else {
-			code.add(", $T.$L, $T.$L)", WebApplicationType.class, webApplicationType,
+		if (!webApplicationType.equals(WebApplicationType.NONE)) {
+			code.add(", $T.$L, $T.$L", WebApplicationType.class, webApplicationType,
 					WebEnvironment.class, detectWebEnvironment(config));
 		}
+		if (args.length > 0) {
+			MultiCodeBlock multi = new MultiCodeBlock();
+			Arrays.stream(args).forEach((arg) -> multi.add("$S", arg));
+			code.add(", $L", multi.join(", "));
+		}
+		code.add(")");
+
 		return code.build();
 	}
 

spring-aot-test/src/main/java/org/springframework/boot/test/context/SpringBootTestArgsAccessor.java

@@ -0,0 +1,38 @@
+/*
+ * Copyright 2019-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.test.context;
+
+import java.util.Set;
+
+import org.springframework.test.context.ContextCustomizer;
+
+/**
+ * Provide access to the package private {@link SpringBootTestArgs} class.
+ *
+ * @author Tadaya Tsuyukubo
+ */
+public class SpringBootTestArgsAccessor {
+
+	public static String[] get(Set<ContextCustomizer> customizers) {
+		return SpringBootTestArgs.get(customizers);
+	}
+
+	public static ContextCustomizer create(Class<?> testClass) {
+		return new SpringBootTestArgs(testClass);
+	}
+
+}

spring-aot-test/src/test/java/org/springframework/aot/test/boot/SpringBootAotTestContextProcessorTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,8 @@
 
 package org.springframework.aot.test.boot;
 
+import java.util.Set;
+
 import com.squareup.javapoet.ClassName;
 import org.junit.jupiter.api.Test;
 
@@ -25,8 +27,11 @@
 import org.springframework.aot.test.samples.app.slice.SampleJdbcTests;
 import org.springframework.boot.WebApplicationType;
 import org.springframework.boot.test.context.ReactiveWebMergedContextConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.context.SpringBootTestArgsAccessor;
 import org.springframework.boot.test.context.SpringBootTestContextBootstrapper;
 import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.test.context.ContextCustomizer;
 import org.springframework.test.context.MergedContextConfiguration;
 import org.springframework.test.context.TestContextBootstrapper;
 import org.springframework.test.context.cache.DefaultCacheAwareContextLoaderDelegate;
@@ -130,10 +135,37 @@ void writeInstanceSupplierForReactiveWebSpringBootTest() {
 				.isEqualTo("() -> new AotSpringBootConfigContextLoader(com.example.Test.class, WebApplicationType.REACTIVE, SpringBootTest.WebEnvironment.MOCK)");
 	}
 
+	@Test
+	void writeInstanceSupplierForNonWebWithArguments() {
+		Set<ContextCustomizer> customizers = Set.of(SpringBootTestArgsAccessor.create(SampleApplicationWithArgumentsTests.class));
+		MergedContextConfiguration contextConfiguration = mock(MergedContextConfiguration.class);
+		given(contextConfiguration.getTestClass()).willAnswer((context) -> SampleApplicationTests.class);
+		given(contextConfiguration.getContextCustomizers()).willAnswer((context) -> customizers);
+		assertThat(CodeSnippet.of(this.processor.writeInstanceSupplier(contextConfiguration, ClassName.get("com.example", "Test"))))
+				.hasImport(AotSpringBootConfigContextLoader.class)
+				.isEqualTo("() -> new AotSpringBootConfigContextLoader(com.example.Test.class, \"--app.test=one\", \"--app.name=foo\")");
+	}
+
+	@Test
+	void writeInstanceSupplierForServletWithArguments() {
+		Set<ContextCustomizer> customizers = Set.of(SpringBootTestArgsAccessor.create(SampleApplicationWithArgumentsTests.class));
+		WebMergedContextConfiguration contextConfiguration = mock(WebMergedContextConfiguration.class);
+		given(contextConfiguration.getTestClass()).willAnswer((context) -> SampleApplicationTests.class);
+		given(contextConfiguration.getContextCustomizers()).willAnswer((context) -> customizers);
+		assertThat(CodeSnippet.of(this.processor.writeInstanceSupplier(contextConfiguration, ClassName.get("com.example", "Test"))))
+				.hasImport(AotSpringBootConfigContextLoader.class).hasImport(WebApplicationType.class)
+				.isEqualTo("() -> new AotSpringBootConfigContextLoader(com.example.Test.class, WebApplicationType.SERVLET, SpringBootTest.WebEnvironment.MOCK, \"--app.test=one\", \"--app.name=foo\")");
+	}
+
 	private TestContextBootstrapper createSpringBootTestContextBootstrapper(Class<?> testClass) {
 		SpringBootTestContextBootstrapper bootstrapper = new SpringBootTestContextBootstrapper();
 		bootstrapper.setBootstrapContext(new DefaultBootstrapContext(testClass, new DefaultCacheAwareContextLoaderDelegate()));
 		return bootstrapper;
 	}
 
+	@SpringBootTest(args = { "--app.test=one", "--app.name=foo" })
+	static class SampleApplicationWithArgumentsTests {
+
+	}
+
 }

spring-native/src/main/java/org/springframework/aot/test/boot/AotSpringBootConfigContextLoader.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -55,34 +55,36 @@ public class AotSpringBootConfigContextLoader extends SpringBootContextLoader {
 
 	private final WebEnvironment webEnvironment;
 
+	private final String[] args;
+
 	/**
 	 * Create an instance for the specified {@link ApplicationContextInitializer} and
 	 * web-related details.
 	 * @param testContextInitializer the context initializer to use
 	 * @param webApplicationType the {@link WebApplicationType} to use for the context
 	 * @param webEnvironment the {@link WebEnvironment} to use for the context
+	 * @param args the command line arguments
 	 */
 	public AotSpringBootConfigContextLoader(Class<? extends ApplicationContextInitializer<?>> testContextInitializer,
-			WebApplicationType webApplicationType, WebEnvironment webEnvironment) {
+			WebApplicationType webApplicationType, WebEnvironment webEnvironment, String... args) {
 		this.testContextInitializer = testContextInitializer;
 		this.webApplicationType = webApplicationType;
 		this.webEnvironment = webEnvironment;
+		this.args = args;
 	}
 
 	/**
 	 * Create a new instance using the specified {@link ApplicationContextInitializer} for
 	 * a non-web context.
 	 * @param testContextInitializer the context initializer to use
+	 * @param args the command line arguments
 	 */
-	public AotSpringBootConfigContextLoader(Class<? extends ApplicationContextInitializer<?>> testContextInitializer) {
-		this(testContextInitializer, WebApplicationType.NONE, WebEnvironment.NONE);
+	public AotSpringBootConfigContextLoader(Class<? extends ApplicationContextInitializer<?>> testContextInitializer, String... args) {
+		this(testContextInitializer, WebApplicationType.NONE, WebEnvironment.NONE, args);
 	}
 
 	@Override
 	public ConfigurableApplicationContext loadContext(MergedContextConfiguration config) {
-		// TODO: handle application arguments
-		String[] args = new String[0];
-
 		SpringApplication application = new AotTestSpringApplication(config.getTestClass().getClassLoader(), testContextInitializer);
 		application.setMainApplicationClass(config.getTestClass());
 		application.setSources(Collections.singleton(testContextInitializer.getName()));
@@ -108,7 +110,7 @@ else if (this.webApplicationType == WebApplicationType.REACTIVE) {
 						ApplicationContextFactory.of(GenericReactiveWebApplicationContext::new));
 			}
 		}
-		ConfigurableApplicationContext context = application.run(args);
+		ConfigurableApplicationContext context = application.run(this.args);
 
 		return context;
 	}

spring-native/src/test/java/org/springframework/aot/test/boot/AotSpringBootConfigContextLoaderTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,12 +16,23 @@
 
 package org.springframework.aot.test.boot;
 
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 
 import org.springframework.beans.factory.support.BeanDefinitionBuilder;
+import org.springframework.boot.ApplicationArguments;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.WebApplicationType;
 import org.springframework.boot.test.autoconfigure.orm.jpa.DataJpaTest;
@@ -84,6 +95,38 @@ void loadContextSetActiveProfiles() {
 				assertThat(context.getEnvironment().getActiveProfiles()).containsOnly("profile1", "profile2"));
 	}
 
+	@ParameterizedTest
+	@MethodSource("applicationArguments")
+	void loadContextUsesArguments(String[] arguments, List<Entry<String, String>> expectedEntries) {
+		AotSpringBootConfigContextLoader loader = new AotSpringBootConfigContextLoader(TestApplicationContextInitializer.class, arguments);
+		run(() -> loader.loadContext(createMergedContextConfiguration(SampleTest.class)), (context) -> {
+			ApplicationArguments args = context.getBean(ApplicationArguments.class);
+			Set<String> keys = expectedEntries.stream().map(Entry::getKey).collect(Collectors.toSet());
+			assertThat(args.getOptionNames()).containsExactlyElementsOf(keys);
+			for (Entry<String, String> expectedEntry : expectedEntries) {
+				String key = expectedEntry.getKey();
+				String value = expectedEntry.getValue();
+				assertThat(args.getOptionValues(key)).containsOnly(value);
+			}
+		});
+	}
+
+	@ParameterizedTest
+	@MethodSource("applicationArguments")
+	void loadContextUsesArgumentsAndWebSettings(String[] arguments, List<Entry<String, String>> expectedEntries) {
+		AotSpringBootConfigContextLoader loader = new AotSpringBootConfigContextLoader(TestApplicationContextInitializer.class, WebApplicationType.SERVLET, WebEnvironment.MOCK, arguments);
+		run(() -> loader.loadContext(createMergedContextConfiguration(SampleTest.class)), (context) -> {
+			ApplicationArguments args = context.getBean(ApplicationArguments.class);
+			Set<String> keys = expectedEntries.stream().map(Entry::getKey).collect(Collectors.toSet());
+			assertThat(args.getOptionNames()).containsExactlyElementsOf(keys);
+			for (Entry<String, String> expectedEntry : expectedEntries) {
+				String key = expectedEntry.getKey();
+				String value = expectedEntry.getValue();
+				assertThat(args.getOptionValues(key)).containsOnly(value);
+			}
+		});
+	}
+
 	private void run(Supplier<ConfigurableApplicationContext> supplier, Consumer<AssertableApplicationContext> context) {
 		try (ConfigurableApplicationContext ctx = supplier.get()) {
 			context.accept(AssertableApplicationContext.get(() -> ctx));
@@ -125,4 +168,12 @@ public void initialize(GenericApplicationContext applicationContext) {
 
 	}
 
+	static Stream<Arguments> applicationArguments() {
+		return Stream.of(
+				Arguments.of(new String[] {}, Collections.emptyList()),
+				Arguments.of(new String[] { "--app.test=one"}, List.of(Map.entry("app.test", "one"))),
+				Arguments.of(new String[] { "--app.test=one", "--app.name=foo" }, List.of(Map.entry("app.test", "one"), Map.entry("app.name", "foo")))
+		);
+	}
+
 }