From ce86b2a295e5fd52fe82676b539466315205325a Mon Sep 17 00:00:00 2001
From: Aleksandr Arshavskiy <aleksandr.arshavskiy@c24.de>
Date: Tue, 17 Jan 2023 14:37:22 +0100
Subject: [PATCH] Disable BatchTestContextCustomizer after AOT processing

---
 .../context/BatchTestContextCustomizer.java   |  8 +++++-
 .../BatchTestContextCustomizerTests.java      | 26 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/spring-batch-test/src/main/java/org/springframework/batch/test/context/BatchTestContextCustomizer.java b/spring-batch-test/src/main/java/org/springframework/batch/test/context/BatchTestContextCustomizer.java
index 7608b88500..955afbe18c 100644
--- a/spring-batch-test/src/main/java/org/springframework/batch/test/context/BatchTestContextCustomizer.java
+++ b/spring-batch-test/src/main/java/org/springframework/batch/test/context/BatchTestContextCustomizer.java
@@ -15,6 +15,7 @@
  */
 package org.springframework.batch.test.context;
 
+import org.springframework.aot.AotDetector;
 import org.springframework.batch.test.JobLauncherTestUtils;
 import org.springframework.batch.test.JobRepositoryTestUtils;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
@@ -30,7 +31,7 @@
  * ({@link JobLauncherTestUtils} and {@link JobRepositoryTestUtils}) as beans in the test
  * context.
  *
- * @author Mahmoud Ben Hassine
+ * @author Mahmoud Ben Hassine, Alexander Arshavskiy
  * @since 4.1
  */
 public class BatchTestContextCustomizer implements ContextCustomizer {
@@ -43,6 +44,11 @@ public class BatchTestContextCustomizer implements ContextCustomizer {
 
 	@Override
 	public void customizeContext(ConfigurableApplicationContext context, MergedContextConfiguration mergedConfig) {
+
+		if (AotDetector.useGeneratedArtifacts()) {
+			return;
+		}
+
 		ConfigurableListableBeanFactory beanFactory = context.getBeanFactory();
 		Assert.isInstanceOf(BeanDefinitionRegistry.class, beanFactory,
 				"The bean factory must be an instance of BeanDefinitionRegistry");
diff --git a/spring-batch-test/src/test/java/org/springframework/batch/test/context/BatchTestContextCustomizerTests.java b/spring-batch-test/src/test/java/org/springframework/batch/test/context/BatchTestContextCustomizerTests.java
index c948d7c10a..e865e2c06f 100644
--- a/spring-batch-test/src/test/java/org/springframework/batch/test/context/BatchTestContextCustomizerTests.java
+++ b/spring-batch-test/src/test/java/org/springframework/batch/test/context/BatchTestContextCustomizerTests.java
@@ -15,25 +15,33 @@
  */
 package org.springframework.batch.test.context;
 
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 
 import org.springframework.context.ConfigurableApplicationContext;
 import org.springframework.context.support.GenericApplicationContext;
+import org.springframework.core.SpringProperties;
 import org.springframework.test.context.MergedContextConfiguration;
 
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 /**
- * @author Mahmoud Ben Hassine
+ * @author Mahmoud Ben Hassine, Alexander Arshavskiy
  */
 class BatchTestContextCustomizerTests {
 
 	private final BatchTestContextCustomizer contextCustomizer = new BatchTestContextCustomizer();
 
+	@AfterEach
+	void removeSystemProperty() {
+		SpringProperties.setProperty("spring.aot.enabled", null);
+	}
+
 	@Test
 	void testCustomizeContext() {
 		// given
@@ -64,4 +72,20 @@ void testCustomizeContext_whenBeanFactoryIsNotAnInstanceOfBeanDefinitionRegistry
 				containsString("The bean factory must be an instance of BeanDefinitionRegistry"));
 	}
 
+	@Test
+	void testCustomizeContext_whenUsingAotGeneratedArtifactsBatchTestContextIsNotRegistered() {
+		// given
+		SpringProperties.setProperty("spring.aot.enabled", "true");
+		ConfigurableApplicationContext context = new GenericApplicationContext();
+		MergedContextConfiguration mergedConfig = Mockito.mock(MergedContextConfiguration.class);
+
+		// when
+		this.contextCustomizer.customizeContext(context, mergedConfig);
+
+		// then
+		assertFalse(context.containsBean("jobLauncherTestUtils"));
+		assertFalse(context.containsBean("jobRepositoryTestUtils"));
+		assertFalse(context.containsBean("batchTestContextBeanPostProcessor"));
+	}
+
 }