#41 fix early initialization of factory beans

Author: tomix26

embedded-database-spring-test/src/main/java/io/zonky/test/db/postgres/EmbeddedPostgresContextCustomizerFactory.java

@@ -40,7 +40,6 @@
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.BeanDefinitionHolder;
 import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
-import org.springframework.beans.factory.config.RuntimeBeanReference;
 import org.springframework.beans.factory.support.BeanDefinitionRegistry;
 import org.springframework.beans.factory.support.BeanDefinitionRegistryPostProcessor;
 import org.springframework.beans.factory.support.RootBeanDefinition;
@@ -235,7 +234,7 @@ public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) t
                     dataSourceDefinition.getConstructorArgumentValues()
                             .addIndexedArgumentValue(0, flywayInfo.getBeanName());
                     dataSourceDefinition.getConstructorArgumentValues()
-                            .addIndexedArgumentValue(1, new RuntimeBeanReference(contextInfo.getBeanName()));
+                            .addIndexedArgumentValue(1, contextInfo.getBeanName());
                 }
 
                 String dataSourceBeanName = dataSourceInfo.getBeanName();

embedded-database-spring-test/src/main/java/io/zonky/test/db/postgres/EmptyEmbeddedPostgresDataSourceFactoryBean.java

@@ -20,29 +20,33 @@
 import io.zonky.test.db.provider.DatabaseDescriptor;
 import io.zonky.test.db.provider.DatabasePreparer;
 import io.zonky.test.db.provider.GenericDatabaseProvider;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.FactoryBean;
-import org.springframework.beans.factory.InitializingBean;
-import org.springframework.beans.factory.annotation.Autowired;
 
 import javax.sql.DataSource;
 
 /**
  * Implementation of the {@link org.springframework.beans.factory.FactoryBean} interface
  * that provides empty instances of the embedded postgres database.
  */
-public class EmptyEmbeddedPostgresDataSourceFactoryBean implements FactoryBean<DataSource>, InitializingBean {
+public class EmptyEmbeddedPostgresDataSourceFactoryBean implements FactoryBean<DataSource>, BeanFactoryAware {
 
     protected final DatabaseDescriptor databaseDescriptor;
 
-    @Autowired
-    protected GenericDatabaseProvider databaseProvider;
-
+    private BeanFactory beanFactory;
     private DataSource dataSource;
 
     public EmptyEmbeddedPostgresDataSourceFactoryBean(DatabaseDescriptor databaseDescriptor) {
         this.databaseDescriptor = databaseDescriptor;
     }
 
+    @Override
+    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
+        this.beanFactory = beanFactory;
+    }
+
     @Override
     public boolean isSingleton() {
         return true;
@@ -54,16 +58,15 @@ public Class<?> getObjectType() {
     }
 
     @Override
-    public DataSource getObject() {
+    public DataSource getObject() throws Exception {
+        if (dataSource == null) {
+            GenericDatabaseProvider databaseProvider = beanFactory.getBean(GenericDatabaseProvider.class);
+            dataSource = databaseProvider.getDatabase(EmptyDatabasePreparer.INSTANCE, databaseDescriptor);
+            EmbeddedDatabaseReporter.reportDataSource(dataSource);
+        }
         return dataSource;
     }
 
-    @Override
-    public void afterPropertiesSet() throws Exception {
-        dataSource = databaseProvider.getDatabase(EmptyDatabasePreparer.INSTANCE, databaseDescriptor);
-        EmbeddedDatabaseReporter.reportDataSource(dataSource);
-    }
-
     private static class EmptyDatabasePreparer implements DatabasePreparer {
 
         public static final EmptyDatabasePreparer INSTANCE = new EmptyDatabasePreparer();

embedded-database-spring-test/src/main/java/io/zonky/test/db/postgres/FlywayEmbeddedPostgresDataSourceFactoryBean.java

@@ -24,8 +24,9 @@
 import org.slf4j.LoggerFactory;
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.BeanFactory;
+import org.springframework.beans.factory.BeanFactoryAware;
 import org.springframework.beans.factory.FactoryBean;
-import org.springframework.beans.factory.InitializingBean;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.core.Ordered;
 import org.springframework.util.concurrent.ListenableFuture;
@@ -37,25 +38,31 @@
  * that provides fully cacheable instances of the embedded postgres database.
  * Each instance is backed by a flyway bean that is used for initializing the target database.
  */
-public class FlywayEmbeddedPostgresDataSourceFactoryBean implements FactoryBean<DataSource>, BeanPostProcessor, InitializingBean, Ordered {
+public class FlywayEmbeddedPostgresDataSourceFactoryBean implements FactoryBean<DataSource>, BeanPostProcessor, BeanFactoryAware, Ordered {
 
     private static final Logger logger = LoggerFactory.getLogger(FlywayEmbeddedPostgresDataSourceFactoryBean.class);
 
     private final String flywayName;
-    private final FlywayDataSourceContext dataSourceContext;
+    private final String dataSourceContextName;
 
+    private BeanFactory beanFactory;
     private DataSource proxyInstance;
 
-    public FlywayEmbeddedPostgresDataSourceFactoryBean(String flywayName, FlywayDataSourceContext dataSourceContext) {
+    public FlywayEmbeddedPostgresDataSourceFactoryBean(String flywayName, String dataSourceContextName) {
         this.flywayName = flywayName;
-        this.dataSourceContext = dataSourceContext;
+        this.dataSourceContextName = dataSourceContextName;
     }
 
     @Override
     public int getOrder() {
         return Ordered.HIGHEST_PRECEDENCE;
     }
 
+    @Override
+    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
+        this.beanFactory = beanFactory;
+    }
+
     @Override
     public boolean isSingleton() {
         return true;
@@ -68,17 +75,17 @@ public Class<?> getObjectType() {
 
     @Override
     public DataSource getObject() {
+        if (proxyInstance == null) {
+            FlywayDataSourceContext dataSourceContext = beanFactory.getBean(dataSourceContextName, FlywayDataSourceContext.class);
+            proxyInstance = ProxyFactory.getProxy(DataSource.class, dataSourceContext);
+        }
         return proxyInstance;
     }
 
-    @Override
-    public void afterPropertiesSet() {
-        proxyInstance = ProxyFactory.getProxy(DataSource.class, dataSourceContext);
-    }
-
     @Override
     public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
         if (bean instanceof Flyway && StringUtils.equals(beanName, flywayName)) {
+            FlywayDataSourceContext dataSourceContext = beanFactory.getBean(dataSourceContextName, FlywayDataSourceContext.class);
             ListenableFuture<DataSource> reloadFuture = dataSourceContext.reload((Flyway) bean);
             reloadFuture.addCallback(
                     result -> EmbeddedDatabaseReporter.reportDataSource(result),

embedded-database-spring-test/src/test/java/io/zonky/test/db/EarlyInitializationIntegrationTest.java

@@ -0,0 +1,65 @@
+package io.zonky.test.db;
+
+import io.zonky.test.category.FlywayIntegrationTests;
+import org.flywaydb.core.Flyway;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.config.BeanFactoryPostProcessor;
+import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.Ordered;
+import org.springframework.test.context.ContextConfiguration;
+import org.springframework.test.context.junit4.SpringRunner;
+
+import javax.sql.DataSource;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+@RunWith(SpringRunner.class)
+@Category(FlywayIntegrationTests.class)
+@AutoConfigureEmbeddedDatabase(beanName = "dataSource")
+@ContextConfiguration
+public class EarlyInitializationIntegrationTest {
+
+    @Configuration
+    static class Config {
+
+        @Bean
+        public BeanFactoryPostProcessor beanFactoryPostProcessor() {
+            return new TestBeanFactoryPostProcessor();
+        }
+
+        @Bean(initMethod = "migrate")
+        public Flyway flyway(DataSource dataSource) {
+            Flyway flyway = new Flyway();
+            flyway.setDataSource(dataSource);
+            flyway.setSchemas("test");
+            return flyway;
+        }
+    }
+
+    @Autowired
+    private DataSource dataSource;
+
+    @Test
+    public void testDataSource() {
+        assertThat(dataSource).isNotNull();
+    }
+
+    private static class TestBeanFactoryPostProcessor implements BeanFactoryPostProcessor, Ordered {
+
+        @Override
+        public int getOrder() {
+            return Ordered.HIGHEST_PRECEDENCE;
+        }
+
+        @Override
+        public void postProcessBeanFactory(ConfigurableListableBeanFactory beanFactory) throws BeansException {
+            beanFactory.getBeanNamesForType(DataSource.class, true, true);
+        }
+    }
+}