AbstractSpringTestNGDBUnitTest: adds DBUnit functionality to Spring's base class for Spring-context aware TestNG tests. To use, subclass and provide XML dataset and DTD files (at locations indicated in API comments).

AbstractSpringTestNGDBUnitTest.java

import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.io.Reader;
import java.sql.SQLException;

import javax.annotation.Resource;

import org.dbunit.DefaultDatabaseTester;
import org.dbunit.IDatabaseTester;
import org.dbunit.IOperationListener;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseDataSet;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.CompositeDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.springframework.core.io.ClassPathResource;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.testng.AbstractTransactionalTestNGSpringContextTests;
import org.springframework.test.context.transaction.TransactionConfiguration;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;

/**
 * <p>Base class for persistence-layer integration tests with DBUnit + Spring + TestNG.
 * Calls DBUnit functions to set up and tear down database before and after each
 * test method, using a test dataset (by default the fully qualified name of the test
 * class followed by ".xml").</p>
 *
 * <p>To ensure complete setup of the database, performs setup and teardown operations
 * not only on all tables in the test dataset, but on all tables listed in the DTD
 * which must be found at the location returned by the getDtdFilename() method.</p>
 *
 * <p>The methods loadFromXml() and getDatabaseDataSet() enable loading datasets from
 * a file and from the database, respectively. This is intended for post-test verification,
 * where you compare the database's actual state with its expected state.  The sister
 * class EqualDataSetAssertBuilder can perform this verification.</p>
 *
 * @author Andrew Spencer
 */
@ContextConfiguration(locations = "/GDP-persistence-test-beans.xml")
@TransactionConfiguration(transactionManager="transactionManagerUdp")
public abstract class AbstractSpringTestNGDBUnitTest extends AbstractTransactionalTestNGSpringContextTests {
    @Resource
    private IDatabaseConnection dbUnitConnection;
    private IDatabaseTester databaseTester;
    private DatabaseOperation setUpOperation;
    private DatabaseOperation tearDownOperation;
    private String xmlFilename;
    private String dtdFilename;
    private Reader dtdReader;

    /**
     * Creates the test with specified parameters.
     * 
     * @param setUpOperation
     *            setup operation to apply before each test method
     * @param tearDownOperation
     *            cleanup operation to apply after each test method
     * @param xmlFilename
     *            location (on classpath) of file containing data (in FlatXmlDataSet format) 
     *            to load before each test
     */
    public AbstractSpringTestNGDBUnitTest(DatabaseOperation setUpOperation, DatabaseOperation tearDownOperation, String xmlFilename) {
        this.setUpOperation = setUpOperation;
        this.tearDownOperation = tearDownOperation;
        this.xmlFilename = xmlFilename;
        this.dtdFilename = getDtdFilename();
    }

    /**
     * Creates the test with default parameters:
     * <ul>
     * <li>setup with CLEAN_INSERT
     * <li>cleanup with NONE
     * <li>test data in "{this.getClass().getName()}.xml"
     * </ul>
     */
    public AbstractSpringTestNGDBUnitTest() {
        this.setUpOperation = DatabaseOperation.CLEAN_INSERT;
        this.tearDownOperation = DatabaseOperation.NONE;
        this.xmlFilename = getClass().getName() + ".xml";
        this.dtdFilename = getDtdFilename();
    }

    /**
     * Creates the test with default values:
     * <ul>
     * <li>setup with CLEAN_INSERT
     * <li>cleanup with NONE
     * </ul>
     * 
     * @param xmlFilename
     *            location (on classpath) of file containing data (in FlatXmlDataSet format) 
     *            to load before each test method
     */
    public AbstractSpringTestNGDBUnitTest(String xmlFilename) {
        this.setUpOperation = DatabaseOperation.CLEAN_INSERT;
        this.tearDownOperation = DatabaseOperation.NONE;
        this.xmlFilename = xmlFilename;
        this.dtdFilename = getDtdFilename();
    }

    /**
     * Returns the location of the DTD used to define the dataset;
     * override this method if the DTD is not at /dbunit.dtd.
     */
    protected String getDtdFilename() {
        return "/dbunit.dtd";
    }

    @BeforeMethod
    public void setUpDatabase() throws Exception {
        DatabaseConfig config = dbUnitConnection.getConfig();
        config.setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new Oracle10DataTypeFactory());
        databaseTester = new DefaultDatabaseTester(dbUnitConnection);
        databaseTester.setOperationListener(IOperationListener.NO_OP_OPERATION_LISTENER);
        databaseTester.setDataSet(loadTestDataWithDtdTableList());
        databaseTester.setSetUpOperation(setUpOperation);
        databaseTester.setTearDownOperation(tearDownOperation);
        databaseTester.onSetup();
    }

    /**
     * Database cleanup operation, annotated to be called after each test method.
     * 
     * @throws Exception
     *             the exception
     */
    @AfterMethod
    public final void tearDownDatabase() throws Exception {
        databaseTester.onTearDown();
        if (dtdReader != null) {
            try {
                dtdReader.close();
            } catch (IOException ignored) {
                //
            }
        }
    }

    /**
     * Loads a dataset into memory from an XML file (FlatXmlDataSet format)
     */
    protected IDataSet loadFromXml(String xmlFileToLoad) throws IOException, DataSetException, SQLException {
        FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
        builder.setMetaDataSet(new DatabaseDataSet(dbUnitConnection, false));
        return builder.build(asFile(xmlFileToLoad));
    }

    private IDataSet loadTestDataWithDtdTableList() throws IOException, DataSetException, SQLException {
        dtdReader = new FileReader(new ClassPathResource(dtdFilename).getFile());
        IDataSet dtdDataset = new FlatDtdDataSet(dtdReader);
        IDataSet xmlDataset = loadFromXml(xmlFilename);
        return new CompositeDataSet(dtdDataset, xmlDataset);
    }

    private File asFile(String filename) throws IOException {
        return new ClassPathResource(filename).getFile();
    }

    /**
     * Loads the database data into memory
     */
    protected final IDataSet getDatabaseDataSet() throws SQLException {
        return dbUnitConnection.createDataSet();
    }
}