Tuesday, February 15, 2011

DbUnit tests with Spring and nested transactions

In my current project we are using DbUnit to set up test data for our system tests. It's a J2EE application that uses Spring and Hibernate. When we started with DbUnit we set it up to have one xml file per test class and configured DbUnit to insert the data in the xml file in a JUnit4 @Before annotated method. But this was kind of slow as the data gets reinserted before each test case. Dbunits practice of commiting its data also meant that even if we extended springs AbstractTransactionalJUnit4SpringContextTests (ATJ4SCT) the dbunit data is not rolled back. Dbunit suggests using the Clean_Insert DatabaseOperation and to have empty table rows in the xml files for tables we dont insert data into but want to have cleared by dbunit. We tried this, but it was cumbersome to always have to put in the empty table rows in the correct order to avoid foreign key constrains on delete. And on a few occasions we got false positives from test cases that passed because an earlier test had inserted data that we had forgotten to remove in the current test.
To get around this we wanted to setup DbUnit to not commit its data, so that it would automatically be rolled back if we used Springs ATJ4SCT. To do this we created a new class extending DbUnits DatabaseDataSourceConnection.

import java.sql.Connection;
import java.sql.SQLException;
import javax.sql.DataSource;
import org.dbunit.database.DatabaseDataSourceConnection;
import org.springframework.jdbc.datasource.DataSourceUtils;
public class SpringDatabaseDataSourceConnection extends DatabaseDataSourceConnection {
private final DataSource dataSource;
public SpringDatabaseDataSourceConnection(DataSource dataSource) throws SQLException {
super(dataSource);
this.dataSource = dataSource;
}
public Connection getConnection() throws SQLException {
Connection conn = DataSourceUtils.getConnection(dataSource);
return new SpringConnection(dataSource, conn);
}
}


And a new class extending java.sql.Connection that delegate all the methods to the passed in connection except for the commit method that we set to do nothing.

public class SpringConnection implements Connection {
private final DataSource dataSource;
private final Connection conn;
public SpringConnection(DataSource dataSource, Connection conn) {
this.dataSource = dataSource;
this.conn = conn;
}
public void commit() throws SQLException {
// Do nothing!
}
/**
* calls DataSourceUtils.closeConnectionIfNecessary rather than directly closing the connection
*/
public void close() throws SQLException {
DataSourceUtils.releaseConnection(conn, dataSource);
}
//Lots of methods here that just delegats to the conn instance variable
}


DataSourceUtils is a Spring class that only closes the DataSource if no one else is using it.
We can use this class in our testcode like so.

DataSource ds = getApplicationContext().getBean("dataSource");
DatabaseDataSourceConnection jdbcConnection = new SpringDatabaseDataSourceConnection(ds);
try {
DatabaseOperation.INSERT.execute(jdbcConnection, createDataSetFromFile("ourTestDataXmlFile.xml"));
finally {
DataSourceUtils.releaseConnection(jdbcConnection.getConnection(), ds);
}


With this setup we could instruct DbUnit to insert the data before each test and not have to worry about putting in empty table statements in our DbUnit xml files. But the test-cases was of course still quite slow to run. It would be nice if we could have DbUnit insert its data before the test class starts and rolled back after all test-cases in the class have finished but still have each test-case run in isolation, in its own transaction.
This is exactly what nested transactions does. It starts a new transaction inside an existing one, that can be rolled back independently of the surrounding transaction but that still sees all data written by the outer transaction. The only problem left was how to start the transaction and insert the test data before each test-class. In a @BeforeClass static method we don't have access to the Spring context. And therefore no access to the Spring configured DataSource or TransactionManager.
But Spring has the notion of TestExecutionListeners, ATJ4SCT uses a TransactionalTestExecutionListener to start and roll back a transaction around each test-method. The TestExecutionListener interface has four methods: beforeTestMethod, afterTestMethod, beforeTestClass and afterTestClass. So by creating our own TestExecutionListener and implement the -testClass methods we will have access to the Spring context and a place to do work before and after each test-class. Our TestExecutionListener looks for a DbUnit xml file with the same name as the test class that is running, to insert as test data. Most of the code is copied from Springs TransactionalTestExecutionListener.

package se.avega.dbunit;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.Map;
import javax.sql.DataSource;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseDataSourceConnection;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.ext.hsqldb.HsqldbDataTypeFactory;
import org.dbunit.operation.DatabaseOperation;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.AbstractTestExecutionListener;
import org.springframework.test.context.transaction.TransactionConfiguration;
import org.springframework.test.context.transaction.TransactionConfigurationAttributes;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.annotation.AnnotationTransactionAttributeSource;
import org.springframework.transaction.interceptor.DefaultTransactionAttribute;
import org.springframework.transaction.interceptor.TransactionAttributeSource;
public class DbUnitTransactionPerTestClassListener extends AbstractTestExecutionListener {
protected final TransactionAttributeSource attributeSource = new AnnotationTransactionAttributeSource();
private TransactionConfigurationAttributes configurationAttributes;
private volatile int transactionsStarted = 0;
private final Map<Class<?>, TransactionContext> transactionContextCache =
Collections.synchronizedMap(new IdentityHashMap<Class<?>, TransactionContext>());
@Override
public void beforeTestClass(TestContext testContext) throws Exception {
final Class<?> testClass = testContext.getTestClass();
if (this.transactionContextCache.remove(testClass) != null) {
throw new IllegalStateException("Cannot start new transaction without ending existing transaction: " +
"Invoke endTransaction() before startNewTransaction().");
}
PlatformTransactionManager tm = getTransactionManager(testContext);
TransactionContext txContext = new TransactionContext(tm,new DefaultTransactionAttribute());
startNewTransaction(testContext, txContext);
this.transactionContextCache.put(testClass, txContext);
insertData(testContext);
}
private void insertData(TestContext testContext) throws Exception {
DataSource ds = (DataSource) testContext.getApplicationContext().getBean("dataSource");
DatabaseDataSourceConnection jdbcConnection = new SpringDatabaseDataSourceConnection(ds);
jdbcConnection.getConfig().setProperty(DatabaseConfig.PROPERTY_DATATYPE_FACTORY, new HsqldbDataTypeFactory());
try {
String filename = testContext.getTestClass().getSimpleName() + ".xml";
DatabaseOperation.INSERT.execute(jdbcConnection, createDataSetFromFile(filename));
} finally {
DataSourceUtils.releaseConnection(jdbcConnection.getConnection(), ds);
}
}
private IDataSet createDataSetFromFile(String fileName) throws DataSetException, FileNotFoundException {
// Our DbUnit xml files is in a directory called testData so add that to the filename.
return new FlatXmlDataSetBuilder().build(new FileInputStream("testData/" + fileName));
}
@Override
public void afterTestClass(TestContext testContext) throws Exception {
Class<?> testClass = testContext.getTestClass();
// If the transaction is still active...
TransactionContext txContext = this.transactionContextCache.remove(testClass);
if (txContext != null && !txContext.transactionStatus.isCompleted()) {
endTransaction(testContext, txContext);
}
}
private void startNewTransaction(TestContext testContext, TransactionContext txContext) throws Exception {
txContext.startTransaction();
++this.transactionsStarted;
}
private void endTransaction(TestContext testContext, TransactionContext txContext) throws Exception {
boolean rollback = isRollback(testContext);
txContext.endTransaction(rollback);
}
protected final PlatformTransactionManager getTransactionManager(TestContext testContext) {
String tmName = retrieveConfigurationAttributes(testContext).getTransactionManagerName();
return testContext.getApplicationContext().getBean(tmName, PlatformTransactionManager.class);
}
protected final boolean isRollback(TestContext testContext) throws Exception {
return true;
}
private TransactionConfigurationAttributes retrieveConfigurationAttributes(TestContext testContext) {
if (this.configurationAttributes == null) {
Class<?> clazz = testContext.getTestClass();
Class<TransactionConfiguration> annotationType = TransactionConfiguration.class;
TransactionConfiguration config = clazz.getAnnotation(annotationType);
String transactionManagerName;
boolean defaultRollback;
if (config != null) {
transactionManagerName = config.transactionManager();
defaultRollback = config.defaultRollback();
}
else {
transactionManagerName = (String) AnnotationUtils.getDefaultValue(annotationType, "transactionManager");
defaultRollback = (Boolean) AnnotationUtils.getDefaultValue(annotationType, "defaultRollback");
}
TransactionConfigurationAttributes configAttributes =
new TransactionConfigurationAttributes(transactionManagerName, defaultRollback);
this.configurationAttributes = configAttributes;
}
return this.configurationAttributes;
}
private static class TransactionContext {
private final PlatformTransactionManager transactionManager;
private final TransactionDefinition transactionDefinition;
private TransactionStatus transactionStatus;
public TransactionContext(PlatformTransactionManager transactionManager, TransactionDefinition transactionDefinition) {
this.transactionManager = transactionManager;
this.transactionDefinition = transactionDefinition;
}
public void startTransaction() {
this.transactionStatus = this.transactionManager.getTransaction(this.transactionDefinition);
}
public void endTransaction(boolean rollback) {
if (rollback) {
this.transactionManager.rollback(this.transactionStatus);
}
else {
this.transactionManager.commit(this.transactionStatus);
}
}
}
}


In our base class for DbUnit test we specified to run the tests with our new TestExecutionListener and to start a nested transaction for each test-case with the Spring @Transactional(propagation=Propagation.NESTED) annotation. The only problem left is that the Hibernate session is not cleared before the outer transaction ends. So in our base class we add a call to clear on the Hibernate session after each test-case.

package se.avega.dbunit;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.TestExecutionListeners;
import org.springframework.test.context.junit4.AbstractTransactionalJUnit4SpringContextTests;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;
@TestExecutionListeners({DbUnitTransactionPerTestClassListener.class})
@ContextConfiguration(locations={"classpath:testApplicationContext.xml"})
@Transactional(propagation=Propagation.NESTED)
public abstract class DbUnitTestCase extends AbstractTransactionalJUnit4SpringContextTests {
@Autowired
private SessionFactory sessionFactory;
@After
public void clearHibernateSession() {
sessionFactory.getCurrentSession().clear();
}
}


Thats it. All integration tests that hits the database can now just extend this base class. If you want to have another naming convention for your DbUnit xml files than the class-name, its easy to create an annotation where you can specify the filename and look for this annotation in the DbUnitTransactionPerTestClassListener class.
Feel free to leave comments about this approach.
How do you handle testdata in your own projects?