diff --git a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java
index 670c871ca..b73f9cc0d 100644
--- a/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java
+++ b/spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/aot/AotFragmentTestConfigurationSupport.java
@@ -57,13 +57,13 @@ class AotFragmentTestConfigurationSupport implements BeanFactoryPostProcessor {
 
 	public AotFragmentTestConfigurationSupport(Class<?> repositoryInterface) {
 		this.repositoryInterface = repositoryInterface;
-		this.repositoryContext = new TestJpaAotRepositoryContext<>(UserRepository.class, null);
+		this.repositoryContext = new TestJpaAotRepositoryContext<>(repositoryInterface, null);
 	}
 
 	@Override
 	public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {
 
-		TestGenerationContext generationContext = new TestGenerationContext(UserRepository.class);
+		TestGenerationContext generationContext = new TestGenerationContext(repositoryInterface);
 
 		new JpaRepositoryContributor(repositoryContext).contribute(generationContext);