com.crosstreelabs.testing.DBunitTester Maven / Gradle / Ivy

Go to download
package com.crosstreelabs.testing;

import com.crosstreelabs.testing.annotation.DBUnit;
import static junit.framework.TestCase.assertNotNull;
import org.dbunit.DataSourceDatabaseTester;
import org.dbunit.DefaultDatabaseTester;
import org.dbunit.DefaultOperationListener;
import org.dbunit.IDatabaseTester;
import org.dbunit.IOperationListener;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.XmlDataSet;
import org.dbunit.operation.DatabaseOperation;
import org.junit.rules.ExternalResource;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * 
 */
public class DBunitTester extends ExternalResource {
    private static final Logger LOG = LoggerFactory.getLogger(DBunitTester.class);
    private IDatabaseTester tester;
    private IOperationListener operationListener;
    private final DataSourceResource datasource;
    private Statement base;
    private Description description;
    
    public DBunitTester(final DataSourceResource datasource) {
        this.datasource = datasource;
    }

    //~ JUnit ExternalResource implementation ~~~~~~~~~~~~~~
    @Override
    public Statement apply(Statement base, Description description) {
        this.base = base;
        this.description = description;
        return super.apply(base, description);
    }

    @Override
    protected void before() throws Throwable {
        LOG.debug("Populate database");
        final IDatabaseTester databaseTester = getDatabaseTester();
        assertNotNull( "DatabaseTester is not set", databaseTester );
        
        String dataSet = "db/data/base.xml";
        if (description.getAnnotation(DBUnit.class) != null) {
            DBUnit dbunit = description.getAnnotation(DBUnit.class);
            if (!dbunit.dataset().trim().isEmpty()) {
                dataSet = dbunit.dataset().trim();
            }
        }
        
        databaseTester.setSetUpOperation( getSetUpOperation() );
        databaseTester.setDataSet( getDataSet(dataSet) );
        databaseTester.setOperationListener( getOperationListener() );
        databaseTester.onSetup();
    }
    
    @Override
    public void after() {
        LOG.debug("Reset database");
        try {
            final IDatabaseTester databaseTester = getDatabaseTester();
            assertNotNull( "DatabaseTester is not set", databaseTester );
        
            String dataSet = "db/data/base.xml";
            if (description.getAnnotation(DBUnit.class) != null) {
                DBUnit dbunit = description.getAnnotation(DBUnit.class);
                if (!dbunit.dataset().trim().isEmpty()) {
                    dataSet = dbunit.dataset().trim();
                }
            }
        
            databaseTester.setTearDownOperation( getTearDownOperation() );
            databaseTester.setDataSet( getDataSet(dataSet) );
            databaseTester.setOperationListener(getOperationListener());
            databaseTester.onTearDown();
        } catch (Exception ex) {
            LOG.error("Exception:",ex);
        } finally {
            tester = null;
        }
    }
    
    //~ DBUnit necessities ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    /**
     * Returns the test database connection.
     * @return 
     * @throws java.lang.Exception
     */
    protected IDatabaseConnection getConnection() throws Exception {
        return newDatabaseTester().getConnection();
    }
    /**
     * Returns the test dataset.
     * @return 
     * @throws java.lang.Exception
     */
    protected IDataSet getDataSet(final String resource) throws Exception {
        return new XmlDataSet(getClass().getClassLoader().getResourceAsStream(resource));
    }

    /**
     * Gets the IDatabaseTester for this testCase.

     * If the IDatabaseTester is not set yet, this method calls
     * newDatabaseTester() to obtain a new instance.
     * @return 
     * @throws Exception
     */
    protected IDatabaseTester getDatabaseTester() throws Exception {
        if(this.tester == null) {
              this.tester = newDatabaseTester();
        }
        return this.tester;
    }

    /**
     * Creates a IDatabaseTester for this testCase.

     *
     * A {@link DefaultDatabaseTester} is used by default.
     * @return
     */
    protected IDatabaseTester newDatabaseTester() {
//        DataSourceDatabaseTester dt = new DataSourceDatabaseTester(datasource.getDataSource(), datasource.getSchema());
//        dt.setOperationListener(getOperationListener());
//        return dt;
        return new DataSourceDatabaseTester(datasource.getDataSource(), datasource.getSchema()){
            @Override
            public IDatabaseConnection getConnection() throws Exception {
                IDatabaseConnection c = super.getConnection();
                c.getConnection().setSchema(getSchema());
                c.getConfig().setProperty(DatabaseConfig.FEATURE_CASE_SENSITIVE_TABLE_NAMES, Boolean.FALSE);
                c.getConfig().setProperty(DatabaseConfig.FEATURE_QUALIFIED_TABLE_NAMES, Boolean.FALSE);
                return c;
            }
        };
    }

    /**
     * Designed to be overridden by subclasses in order to set additional configuration
     * parameters for the {@link IDatabaseConnection}.
     * @param config The settings of the current {@link IDatabaseConnection} to be configured
     */
    protected void setUpDatabaseConfig(DatabaseConfig config) { }

    /**
     * Returns the database operation executed in test setup.
     * @return 
     * @throws java.lang.Exception 
     */
    protected DatabaseOperation getSetUpOperation() throws Exception {
        return DatabaseOperation.CLEAN_INSERT;
    }

    /**
     * Returns the database operation executed in test cleanup.
     * @return 
     * @throws java.lang.Exception 
     */
    protected DatabaseOperation getTearDownOperation() throws Exception {
        return DatabaseOperation.NONE;
    }
    /**
     * @return The {@link IOperationListener} to be used by the {@link IDatabaseTester}.
     * @since 2.4.4
     */
    protected IOperationListener getOperationListener() {
        LOG.debug("getOperationListener() - start");
        if(this.operationListener==null) {
            this.operationListener = new DefaultOperationListener() {
                @Override
                public void connectionRetrieved(IDatabaseConnection connection) {
                    super.connectionRetrieved(connection);
                    // When a new connection has been created then invoke the setUp method
                    // so that user defined DatabaseConfig parameters can be set.
                    setUpDatabaseConfig(connection.getConfig());
                }
            };
        }
        return this.operationListener;
    }
}