解决unitils中的dbunit使用spring中定义的多数据源

最近在研究unitils, dbunit来适应目前的单元测试.
在unitils中将要使用的数据源都定义在unitils.properties中, 而在我们的测试配置中, 同样的数据源也在spring中配置了一份儿, 因此本人希望同样的配置不应该出现在两个方面, 从而增加维护成本. 另一方面还可以通过spring来解决多数据源的问题. 于是开始看unitils的源代码, 原来在dbunit module中, 获取数据源的代码:

public DbUnitDatabaseConnection getDbUnitDatabaseConnection(String schemaName) {
        DbUnitDatabaseConnection dbUnitDatabaseConnection = dbUnitDatabaseConnections.get(schemaName);
        if (dbUnitDatabaseConnection == null) {
            dbUnitDatabaseConnection = createDbUnitConnection(schemaName);
            dbUnitDatabaseConnections.put(schemaName, dbUnitDatabaseConnection);
        }
        return dbUnitDatabaseConnection;
    }
    protected DbUnitDatabaseConnection createDbUnitConnection(String schemaName) {
        // A DbSupport instance is fetched in order to get the schema name in correct case
        DataSource dataSource = getDatabaseModule().getDataSourceAndActivateTransactionIfNeeded();
        ...
    }

也就是说, 最终的数据源是通过DatabaseModule来获取的, 因此我们需要对这些代码进行改造, 但是还有另外一个问题, unitils的spring module对所有测试的spring context都是按照以测试类为key, context为value进行缓存的. 因此要获取spring的context, 必须知道当前的测试类, 因此这里需要整体调整Dbunit module的insertDataset()方法的内部方法调用的参数. 在处理流程中增加testObject参数. 这个需要修改的代码非常多, 这里不一一展开, 对于从spring中获取数据源的代码如下:

    private DataSource getDataSource(String schemaName, Object testObject) {
        // 从spring中取
        SpringModule springModule = Unitils.getInstance().getModulesRepository().getModuleOfType(SpringModule.class);
        DataSource dataSource = (DataSource) springModule.getSpringBean(testObject, schemaName);
        if (dataSource == null) {
            throw new RuntimeException(String.format("datasource[%s]在spring配置文件中不存在", schemaName));
        }
        return dataSource;
    }

为了将spring的datasource与dbunit能够关联起来, 本人对构造的测试数据做了一个约定, 对于操作的数据表, 必须加数据源标识前缀, 而数据源标识则对应spring bean配置文件中的id, 比如配置了一个datasource bean:

	<bean id="db1" parent="parentDataSource">
		<property name="url">
			<value>jdbc:oracle:oci:@${db1.name}</value>
		</property>
		<property name="username">
			<value>${db1.username}</value>
		</property>
		<property name="password">
			<value>${db1.password}</value>
		</property>
	</bean>

如果有一个dataset中有一个table(table1)跟该datasource关系, 那么该表名必须这样定义: db1.table1. 也可以将db1前缀理解为schema(但这里实际并不是).对于这个问题, 已经跟unitils的founder进行了交流, 并在jira中增加了一个issue: http://jira.unitils.org/browse/UNI-190 在未来版本中将实现该功能.

在对数据清理的处理上的改进
对于构造的测试准备数据的清理, dbunit在默认情况下, 假定所有测试的数据库表必须建立主键, 否则会抛出异常, 因为在清理数据的时候需要利用主键以及构造的对应主键值来做db的delete操作.而我们目前存在一些关联表是没有主键的, 因此给使用unitils带来了一定的麻烦. 但也是可以搞定的. 另外一个问题就是, 对于所有的构造数据都会根据dataset中定义的table来处理, 而目前我们面临的另外一个问题, 就是在测试前, 希望能通过指定的sql来对测试环境对某些数据进行清理操作, 以避免对测试造成影响. 于是对unitils的DataSetLoadStrategy和dbunit的DatabaseOperation进行了扩展.
本人定义了一个ExecuteSqlOperation用来执行指定的sql语句:

public class ExecuteSqlOperation extends DatabaseOperation {
    private Map<String, List<String>> sqlMap;
    private static final Logger logger = LoggerFactory.getLogger(ExecuteSqlOperation.class);

    public ExecuteSqlOperation(Map<String, List<String>> sqlMap) {
        this.sqlMap = sqlMap;
    }

    @Override
    public void execute(IDatabaseConnection connection, IDataSet dataSet) throws DatabaseUnitException, SQLException {
        logger.debug("execute(connection={}, dataSet={}) - start", connection, dataSet);

        DatabaseConfig databaseConfig = connection.getConfig();
        IStatementFactory factory =
                (IStatementFactory) databaseConfig.getProperty(DatabaseConfig.PROPERTY_STATEMENT_FACTORY);

        // for each table
        Iterator<Entry<String, List<String>>> iterator = sqlMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Entry<String, List<String>> sqlEntry = iterator.next();
            String schema = sqlEntry.getKey();

            // Do not process empty table
            if (StringUtils.isBlank(schema) || sqlEntry.getValue().isEmpty()) {
                logger.warn(String.format("schema{%s} or sqls(%s) is empty", schema, sqlEntry.getValue()));
                continue;
            }

            // don't find the datasource
            if (!connection.getSchema().equalsIgnoreCase(schema)) {
                continue;
            }

            IPreparedBatchStatement statement = null;

            try {
                for (String sql : sqlEntry.getValue()) {
                    statement = factory.createPreparedBatchStatement(sql, connection);
                    statement.addBatch();
                }

                statement.executeBatch();
                statement.clearBatch();

                // clear schema and sql
                iterator.remove();
            } finally {
                if (statement != null) {
                    statement.close();
                }
            }
        }
    }

    public void setSqlMap(Map<String, List<String>> sqlMap) {
        this.sqlMap = sqlMap;
    }
}

从代码中我们可以看出, dbunit中的DatabaseOperation 与dataset是有非常强的耦合的(本来嘛, dataset是整个dbunit的模型核心), 而这里我们对dataset是不需要的: 拿到connection, 执行sql, 退出.
然后定义了一个CustemAndInsertDataSetLoadStrategy, 用来引入我们上面定义的这个ExecuteSqlOperation :

public abstract class CustemAndInsertDataSetLoadStrategy implements DataSetLoadStrategy {
    public void execute(DbUnitDatabaseConnection dbUnitDatabaseConnection, IDataSet dataSet) {
        try {
            doExecute(dbUnitDatabaseConnection, dataSet);
            DatabaseOperation.INSERT.execute(dbUnitDatabaseConnection, dataSet);
        } catch (DatabaseUnitException e) {
            throw new UnitilsException("Error while executing DataSetLoadStrategy", e);
        } catch (SQLException e) {
            throw new UnitilsException("Error while executing DataSetLoadStrategy", e);
        }
    }

    abstract protected void doExecute(DbUnitDatabaseConnection dbUnitDatabaseConnection, IDataSet dataSet)
            throws DatabaseUnitException, SQLException;
}

在实际测试中使用:

@DataSet(loadStrategy = MyTest.MyDataSetLoadStrategy.class)
public class MyTest extends IcBaseCase2 {
    public static class MyDataSetLoadStrategy extends CustemAndInsertDataSetLoadStrategy {
        private Map<String, List<String>> sqlMap = new HashMap<String, List<String>>();
        private ExecuteSqlOperation executeSqlOperation;
        public SpuRelationDataSetLoadStrategy() {
            sqlMap.put("db1", Collections
                    .singletonList("delete table1 where id = 110"));
            executeSqlOperation = new ExecuteSqlOperation(sqlMap);
        }

        @Override
        protected void doExecute(DbUnitDatabaseConnection dbUnitDatabaseConnection, IDataSet dataSet)
                throws DatabaseUnitException, SQLException {
            executeSqlOperation.execute(dbUnitDatabaseConnection, dataSet);
        }
    }

当然中间省略了一些实现细节, 发现经过这样扩展之后, 基本能满足我们的测试需要.
感觉多数据源的使用是一个比较少的场景, 从dbunit, unitils的feature中就可以看出来, 貌似二者基本上没有做实现, 因此有了上面的代码.