org.reusables.dbunit.DbUnitDatasetExecutionListener.java Source code

Java tutorial

Introduction

Here is the source code for org.reusables.dbunit.DbUnitDatasetExecutionListener.java

Source

/*
 * Copyright 2010-2012 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.reusables.dbunit;

import static java.lang.String.format;
import static org.springframework.core.annotation.AnnotationUtils.findAnnotation;

import org.reusables.dbunit.autocomplete.AutoCompletionDataset;
import org.reusables.dbunit.handler.DataSourceOperationHandler;
import org.reusables.dbunit.handler.HibernateEntityManagerOperationHandler;
import org.reusables.dbunit.handler.HibernateSessionFactoryOperationHandler;
import org.reusables.dbunit.handler.JdbcTask;

import java.io.InputStream;
import java.net.URL;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collection;

import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang.ArrayUtils;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.dataset.CompositeDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.dataset.xml.XmlDataSet;
import org.dbunit.operation.DatabaseOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.AbstractTestExecutionListener;

/**
 * Checks if the test class or test method is annotated with {@link DbUnitDataset}. If this is the case the
 * given dataset is loaded using DbUnit.
 * 
 * <p>Example test class:
 * </p>
 * 
 * <pre>
 * &#064;TestExecutionListeners({DbUnitDatasetExecutionListener.class})
 * &#064;ContextConfiguration(locations = "classpath:/testContext.xml")
 * &#064;DbUnitDataset("/data/PersonRepositoryTest.xml")
 * public class PersonRepositoryTest extends AbstractTransactionalJUnit4SpringContextTests
 * {
 * ...
 * </pre>
 * 
 * <p>This listener can also use the auto completion feature, by specifying the {@link DbUnitCompletionRules} annotation on your test class,
 * with the file containing the dataset auto completion rules, e.g.:
 * </p>
 * 
 * <pre>
 * &#064;TestExecutionListeners({DbUnitDatasetExecutionListener.class})
 * &#064;DbUnitDataset("/data/PersonRepositoryTest.xml")
 * &#064;DbUnitCompletionRules("/dataset-rules.xml")
 * public class PersonRepositoryTest extends AbstractTransactionalJUnit4SpringContextTests
 * {
 * ...
 * </pre>
 * 
 * <p>You can also omit the {@link DbUnitDataset} annotation. The listener then scans for a dataset file on
 * default locations in the classpath.
 * </p>
 * 
 * <p>For all methods in the test class these locations are scanned, of which the first match is applied:
 * </p>
 * <ol>
 * <li><b>/data/&lt;test-class-name&gt;/class.xml</b> (example: '/data/MyServiceTest/class.xml')</li>
 * <li><b>/data/&lt;test-class-name&gt;/&lt;test-class-name&gt;.xml</b> (example: '/data/MyServiceTest/MyServiceTest.xml')</li>
 * <li><b>/data/&lt;test-class-name&gt;.xml</b> (example: '/data/MyServiceTest.xml')</li>
 * </ol>
 * 
 * <p>In addition these locations are scanned for the current test method (first match is applied):
 * </p>
 * <ol>
 * <li><b>/data/&lt;test-class-name&gt;/&lt;test-method-name&gt;.xml</b> (example: '/data/MyServiceTest/testMethod.xml')</li>
 * <li><b>/data/&lt;test-class-name&gt;_&lt;test-method-name&gt;.xml</b> (example: '/data/MyServiceTest_testMethod.xml')</li>
 * <li><b>/data/&lt;test-class-name&gt;-&lt;test-method-name&gt;.xml</b> (example: '/data/MyServiceTest-testMethod.xml')</li>
 * </ol>
 * 
 * <p>When automatically determining the dataset locations, each test method can have one class level dataset and one method level dataset.
 * </p>
 * 
 * <p>Using the {@link DbUnitBehavior} annotation you can change the behavior of the listener, like the DbUnit operation used for inserting the
 * dataset before the test and cleaning up the dataset after the test. You can also change the way the listener obtains it's datasource, e.g.
 * borrow the dataset of a Hibernate session.
 * </p>
 * 
 * @author marcel
 * @see AutoCompletionDataset
 * @see DbUnitDataset
 * @see DbUnitBehavior
 **/
public class DbUnitDatasetExecutionListener extends AbstractTestExecutionListener {
    private static final Logger LOG = LoggerFactory.getLogger(DbUnitDatasetExecutionListener.class);

    @Override
    public void beforeTestMethod(final TestContext testContext) throws Exception {
        final SetupOperation operation = getInsertOperation(testContext,
                findAnnotation(testContext.getTestClass(), DbUnitDataset.class));

        handleDataset(operation.getOperation(), testContext,
                findAnnotation(testContext.getTestClass(), DbUnitDataset.class),
                findAnnotation(testContext.getTestMethod(), DbUnitDataset.class));
    }

    @Override
    public void afterTestMethod(final TestContext testContext) throws Exception {
        final DbUnitDataset classAnnotation = findAnnotation(testContext.getTestClass(), DbUnitDataset.class);
        final TeardownOperation operation = getCleanAfterOperation(testContext, classAnnotation);

        if (operation != null && operation != TeardownOperation.NONE) {
            final DbUnitDataset methodAnnotation = findAnnotation(testContext.getTestMethod(), DbUnitDataset.class);
            handleDataset(operation.getOperation(), testContext, classAnnotation, methodAnnotation);
        }
    }

    /**
     * Load the dataset using the given annotation information.
     * 
     * @param operation The operation to perform.
     * @param testContext The context for the current test.
     * @param classAnnotation Dataset resource information of the test class.
     * @param methodAnnotation Dataset resource information of the test method.
     * @throws Exception Any error.
     * @since 1.3.0
     */
    protected void handleDataset(final DatabaseOperation operation, final TestContext testContext,
            final DbUnitDataset classAnnotation, final DbUnitDataset methodAnnotation) throws Exception {
        LOG.debug("Getting datasets for operation: {}", operation.getClass().getSimpleName());

        final Collection<IDataSet> datasets = new ArrayList<IDataSet>();
        datasets.addAll(createDataSets(testContext, classAnnotation, false));
        datasets.addAll(createDataSets(testContext, methodAnnotation, true));

        if (datasets.isEmpty()) {
            LOG.debug("No datasets found.");
            return;
        }

        final IDataSet dataset = createAutoCompletionDataSet(
                new CompositeDataSet(datasets.toArray(new IDataSet[0])), testContext);
        final ConnectionType connectionType = getConnectionType(testContext, classAnnotation);
        handleOperation(new DatasetTask(operation, dataset), testContext, connectionType);
    }

    /**
     * Create the datasets for the given annotation. If no annotation is supplied
     * or when the annotation contains no dataset resources, this method tries to find a dataset on the default location.
     * 
     * @param testContext The context for the current test.
     * @param annotation The dataset resource information.
     * @param methodAnnotation If true, indicates that the annotation being processed is a method level annotation, false is for class level annotations.
     * @return The datasets found for the given annotation, or default dataset.
     * 
     * @throws Exception Any error.
     * @since 1.3.0
     */
    protected Collection<IDataSet> createDataSets(final TestContext testContext, final DbUnitDataset annotation,
            final boolean methodAnnotation) throws Exception {
        final Collection<IDataSet> datasets = new ArrayList<IDataSet>();

        if (annotation == null || ArrayUtils.isEmpty(annotation.value())) {
            if (methodAnnotation) {
                CollectionUtils.addIgnoreNull(datasets, getDefaultMethodDataSet(testContext, annotation));
            } else {
                CollectionUtils.addIgnoreNull(datasets, getDefaultClassDataSet(testContext, annotation));
            }
        } else {
            for (final String resource : annotation.value()) {
                LOG.debug("Adding dataset file from classpath '{}'", resource);
                datasets.add(createDataSet(getClass().getResourceAsStream(resource), annotation));
            }
        }
        return datasets;
    }

    /**
     * Try to find the dataset for the test class, by trying a number of default locations.
     * @param testContext The context for the current test.
     * @param annotation Optional - the dataset annotation on the class.
     * 
     * @return The dataset found or null.
     * @throws Exception Any error.
     * @since 1.3.0
     */
    protected IDataSet getDefaultClassDataSet(final TestContext testContext, final DbUnitDataset annotation)
            throws Exception {
        final String className = testContext.getTestClass().getSimpleName();
        IDataSet dataset = getDataSet(format("/data/%s/class.xml", className), annotation);
        if (dataset != null) {
            return dataset;
        }

        dataset = getDataSet(format("/data/%s/%s.xml", className, className), annotation);
        if (dataset != null) {
            return dataset;
        }

        return getDataSet(format("/data/%s.xml", className), annotation);
    }

    /**
     * Try to find the dataset for the test method, by trying a number of default locations.
     * @param testContext The context for the current test.
     * @param annotation Optional - the dataset annotation on the method.
     * 
     * @return The dataset found or null.
     * @throws Exception Any error.
     * @since 1.3.0
     */
    protected IDataSet getDefaultMethodDataSet(final TestContext testContext, final DbUnitDataset annotation)
            throws Exception {
        final String className = testContext.getTestClass().getSimpleName();
        final String methodName = testContext.getTestMethod().getName();
        IDataSet dataset = getDataSet(format("/data/%s/%s.xml", className, methodName), annotation);
        if (dataset != null) {
            return dataset;
        }

        dataset = getDataSet(format("/data/%s_%s.xml", className, methodName), annotation);
        if (dataset != null) {
            return dataset;
        }

        return getDataSet(format("/data/%s-%s.xml", className, methodName), annotation);
    }

    /**
     * Try loading the dataset for the given resource.
     * @param resource The resource to try to load.
     * @param annotation Optional - the dataset annotation.
     * 
     * @return The dataset found or null.
     * @throws Exception
     * @since 1.3.0
     */
    protected IDataSet getDataSet(final String resource, final DbUnitDataset annotation) throws Exception {
        final InputStream input = getClass().getResourceAsStream(resource);

        if (input != null) {
            LOG.debug("Creating dataset for classpath resource: '{}'", resource);
            return createDataSet(input, annotation);
        }

        return null;
    }

    /**
     * Create a dataset for the given inputstream.
     * 
     * @param inputStream The inputsteam to use for the dataset.
     * @param annotation Resource information.
     * @return The dataset.
     * @throws Exception On any error.
     */
    protected IDataSet createDataSet(final InputStream inputStream, final DbUnitDataset annotation)
            throws Exception {
        if (annotation != null && DatasetType.XML.equals(annotation.type())) {
            return new XmlDataSet(inputStream);
        }

        final FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
        builder.setColumnSensing(annotation == null ? true : annotation.columnSensing());
        return builder.build(inputStream);
    }

    /**
     * Determain the connection type to use. If the classAnnotation defines a connection type, this type is returned.
     * Else the default connection type is returned, which is {@link ConnectionType#AUTO}.
     * 
     * @param testContext Listeren test context.
     * @param classAnnotation Class level annotation for behavior information.
     * 
     * @return The connection type to use.
     * @throws Exception On any error.
     * @since 1.3.0
     * @see DbUnitBehavior#connectionType()
     */
    protected ConnectionType getConnectionType(final TestContext testContext, final DbUnitDataset classAnnotation)
            throws Exception {
        final DbUnitBehavior behavior = getBehavior(testContext);
        return behavior != null ? behavior.connectionType() : ConnectionType.AUTO;
    }

    /**
     * Load the testdata from the given inputstream.
     * 
     * @param task Task for handling the DbUnit operation.
     * @param testContext The context for the current test.
     * @param connectionType Type indicating how to obtain a connection.
     * @throws Exception On any error.
     */
    protected void handleOperation(final JdbcTask task, final TestContext testContext,
            final ConnectionType connectionType) throws Exception {
        switch (connectionType) {
        case HIBERNATE_ENTITY_MANAGER:
            if (!new HibernateEntityManagerOperationHandler(testContext).handleOperation(task)) {
                throw new IllegalAccessException("Invalid connectionType: " + connectionType);
            }
            break;
        case HIBERNATE_SESSION_FACTORY:
            if (!new HibernateSessionFactoryOperationHandler(testContext).handleOperation(task)) {
                throw new IllegalAccessException("Invalid connectionType: " + connectionType);
            }
            break;
        case DATA_SOURCE:
            if (!new DataSourceOperationHandler(testContext).handleOperation(task)) {
                throw new IllegalAccessException("Invalid connectionType: " + connectionType);
            }
            break;
        case AUTO:
            if (!(new HibernateEntityManagerOperationHandler(testContext).handleOperation(task)
                    || new HibernateSessionFactoryOperationHandler(testContext).handleOperation(task)
                    || new DataSourceOperationHandler(testContext).handleOperation(task))) {
                throw new IllegalArgumentException("Unable to execute operation, no handler available.");
            }
            break;
        }
    }

    /**
     * Execute the given operation for the given test dataset using the given connection.
     * @param operation DbUnit operation.
     * @param connection Database connection.
     * @param dataset The data set.
     */
    protected void handleOperation(final DatabaseOperation operation, final Connection connection,
            final IDataSet dataset) {
        try {
            operation.execute(new DatabaseConnection(connection), dataset);
        } catch (final DatabaseUnitException e) {
            throw new DbUnitException("Error executing operation.", e);
        } catch (final SQLException e) {
            throw new DbUnitException("SQL error executing operation.", e);
        }
    }

    /**
     * Get the operation to use for setting up the database content.
     * <p>
     * When {@link DbUnitBehavior#setupOperation()} is specified, this is used as setup operation,
     * otherwise {@link SetupOperation#REFRESH} is used.
     * </p>
     * 
     * @param testContext Listener test context.
     * @param classAnnotation Class level annotation for behavior information.
     * @return The operation to use.
     * @since 1.3.0
     */
    protected SetupOperation getInsertOperation(final TestContext testContext,
            final DbUnitDataset classAnnotation) {
        final DbUnitBehavior behavior = getBehavior(testContext);
        return behavior != null ? behavior.setupOperation() : SetupOperation.REFRESH;
    }

    /**
     * Check if the tables of the datasets must be cleaned up after the test, using a specified teardown operation using
     * {@link DbUnitBehavior#teardownOperation()}.
     * 
     * @param testContext The context for the current test.
     * @param classAnnotation Dataset resource information of the test class.
     * 
     * @return The opertation to use for cleaning up after the test, or {@link TeardownOperation#NONE} for no operation.
     */
    protected TeardownOperation getCleanAfterOperation(final TestContext testContext,
            final DbUnitDataset classAnnotation) {
        final DbUnitBehavior behavior = getBehavior(testContext);
        return behavior != null ? behavior.teardownOperation() : TeardownOperation.NONE;
    }

    /**
     * Called to allow the given dataset to be wrapped when auto completion applies.
     * 
     * @param targetDataSet The original dataset.
     * @param testContext The test context.
     * @return The auto completion dataset or if no auto completion is configured, the targetDataSet.
     * @throws Exception Doesn't thow any checked exceptions, but overriding implementations are allowed to do so.
     * @throws DbUnitException When the resource specified via {@link DbUnitCompletionRules} does not exist.
     * @since 1.3.0
     * @see AutoCompletionDataset
     */
    protected IDataSet createAutoCompletionDataSet(final IDataSet targetDataSet, final TestContext testContext)
            throws Exception {
        final URL autoCompleteRulesUrl = getAutoCompletionRulesUrl(testContext);

        if (autoCompleteRulesUrl != null) {
            return new AutoCompletionDataset(targetDataSet, autoCompleteRulesUrl);
        }

        return targetDataSet;
    }

    /**
     * 
     * Checks if auto completion is to be used and returns the URL for the auto completion rules configuration.
     * 
     * @param testContext The test context.
     * @return The auto completion rules file URL or null when not applicable.
     * @throws Exception Doesn't thow any checked exceptions, but overriding implementations are allowed to do so.
     * @throws DbUnitException When the resource specified via {@link DbUnitCompletionRules} does not exist.
     * @since 1.3.0
     */
    protected URL getAutoCompletionRulesUrl(final TestContext testContext) throws Exception {
        final DbUnitCompletionRules annotation = findAnnotation(testContext.getTestClass(),
                DbUnitCompletionRules.class);

        if (annotation == null) {
            return null;
        }

        final URL resource = getClass().getResource(annotation.value());
        if (resource == null) {
            throw new DbUnitException("Resource not found in classpath: " + annotation.value());
        }

        return resource;
    }

    private DbUnitBehavior getBehavior(final TestContext testContext) {
        return findAnnotation(testContext.getTestClass(), DbUnitBehavior.class);
    }

    private final class DatasetTask implements JdbcTask {
        private final DatabaseOperation operation;
        private final IDataSet dataset;

        DatasetTask(final DatabaseOperation operation, final IDataSet dataset) {
            this.operation = operation;
            this.dataset = dataset;
        }

        @Override
        public void execute(final Connection connection) throws SQLException {
            handleOperation(this.operation, connection, this.dataset);
        }
    }
}