H2Rule.java

H2Rule.java

package junit.rules;

import org.h2.jdbcx.JdbcConnectionPool;
import org.junit.rules.ExternalResource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.function.Consumer;

public class H2Rule extends ExternalResource {

    private final static Logger log = LoggerFactory.getLogger(H2Rule.class);


    private final String dbName;

    private JdbcConnectionPool connectionPool;
    private Connection keepAliveConnection;

    public H2Rule(String dbName) {
        this.dbName = dbName;
    }

    @Override
    protected void before() throws Throwable {
        connectionPool = JdbcConnectionPool.create("jdbc:h2:mem:" + dbName, "sa", "sa"); // <1>
        keepAliveConnection = connectionPool.getConnection(); // <2>
    }

    public Connection openConnection() throws SQLException {
        return connectionPool.getConnection();
    }

    public void withConnection(Consumer<Connection> consumer) throws SQLException {
        try (Connection connection = openConnection()) {
            consumer.accept(connection);
        }
    }

    // <3>

    @Override
    protected void after() {
        if (keepAliveConnection != null) {
            try {
                keepAliveConnection.close();
            }
            catch (SQLException ex) {
                log.warn("Failed to close connection used to keep H2 database open during test execution", ex);
            }
        }

        if (connectionPool != null) {
            connectionPool.dispose();
        }
    }

}