DatabaseConfiguratorExtension.java

import java.io.IOException;
import java.sql.SQLException;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Consumer;

import javax.sql.DataSource;

import org.flywaydb.core.Flyway;
import org.junit.jupiter.api.extension.AfterAllCallback;
import org.junit.jupiter.api.extension.BeforeAllCallback;
import org.junit.jupiter.api.extension.ExtensionContext;
import org.junit.jupiter.api.extension.ParameterContext;
import org.junit.jupiter.api.extension.ParameterResolutionException;
import org.junit.jupiter.api.extension.ParameterResolver;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;

import io.zonky.test.db.postgres.embedded.EmbeddedPostgres;

public class DatabaseConfiguratorExtension implements BeforeAllCallback, AfterAllCallback, ParameterResolver {

  public static final String JDBC_TEMPLATE_KEY = "NamedParameterJdbcTemplate";
  public static final String DATABASE_CONFIGURATOR_KEY = "DatabaseConfigurator";
  public static final String DATA_SOURCE_KEY = "DataSource";
  private final List<Consumer<EmbeddedPostgres.Builder>> builderCustomizers = new CopyOnWriteArrayList<>();

  public DatabaseConfiguratorExtension() {
  }

  @Override
  public void beforeAll(ExtensionContext context) throws Exception {

    final ExtensionContext.Store store = this.getStore(context);

    if (this.getDataSource(store) == null) {
      setNewInstanceOnContext(store);
    }

  }

  @Override
  public void afterAll(ExtensionContext context) throws Exception {
    this.getStore(context)
        .get("instance", EmbeddedPostgres.class)
        .close();
  }

  private ExtensionContext.Store getStore(ExtensionContext context) {
    return context
        .getRoot()
        .getStore(ExtensionContext.Namespace.GLOBAL);
  }

  private void setNewInstanceOnContext(ExtensionContext.Store store) throws SQLException, IOException {
    this.customize(customizer -> {
      customizer.setPort(5430);
    });
    final EmbeddedPostgres instance = this.pg();
    final DataSource dc = instance.getPostgresDatabase();
    store.put("instance", instance);
    store.put(JDBC_TEMPLATE_KEY, new NamedParameterJdbcTemplate(dc));
    store.put(DATA_SOURCE_KEY, dc);
    store.put(DATABASE_CONFIGURATOR_KEY, new DatabaseConfigurator(new NamedParameterJdbcTemplate(dc)));
  }

  private EmbeddedPostgres pg() throws IOException {
    final EmbeddedPostgres.Builder builder = EmbeddedPostgres.builder();
    this.builderCustomizers.forEach(it -> it.accept(builder));
    final EmbeddedPostgres embeddedPostgres = builder.start();
    final Flyway flyway = new Flyway(Flyway
        .configure()
        .dataSource(embeddedPostgres.getPostgresDatabase())
        .locations("db/migrations")
    );
    flyway.migrate();
    return embeddedPostgres;
  }

  public DatabaseConfiguratorExtension customize(Consumer<EmbeddedPostgres.Builder> customizer) {
    this.builderCustomizers.add(customizer);
    return this;
  }

  @Override
  public boolean supportsParameter(ParameterContext parameterContext, ExtensionContext extensionContext)
      throws ParameterResolutionException {
    return this.resolveParameter(parameterContext, extensionContext) != null;
  }

  @Override
  public Object resolveParameter(ParameterContext parameterContext, ExtensionContext extensionContext)
      throws ParameterResolutionException {
    switch (parameterContext.getParameter()
        .getType()
        .getSimpleName()) {
      case DATA_SOURCE_KEY:
        return this.getDataSource(this.getStore(extensionContext));
      case JDBC_TEMPLATE_KEY:
        return this.getJdbcTemplate(extensionContext);
      case DATABASE_CONFIGURATOR_KEY:
        return this.getDatabaseConfigurator(extensionContext);
      default:
        return null;
    }
  }

  private DatabaseConfigurator getDatabaseConfigurator(ExtensionContext extensionContext) {
    return this.getStore(extensionContext)
        .get(DATABASE_CONFIGURATOR_KEY, DatabaseConfigurator.class);
  }

  private NamedParameterJdbcTemplate getJdbcTemplate(ExtensionContext extensionContext) {
    return this.getStore(extensionContext)
        .get(JDBC_TEMPLATE_KEY, NamedParameterJdbcTemplate.class);
  }

  private DataSource getDataSource(ExtensionContext.Store store) {
    return store.get(DATA_SOURCE_KEY, DataSource.class);
  }

}