Dynamic Replacers for DBUnit

DynamicReplacementTable.java

package io.github.tourniquet.replacer


import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ReplacementTable;

import io.github.tourniquet.replacer.DynamicReplacementDataSet.Replacer;


/**
 * Extends {@link ReplacementTable} with the functionality to replace regexes instead of only strings by adding
 * {@link Replacer}s.
 * Replacements added via {@link #addReplacementObject(Object, Object)} have precedence over any Replacers.
 * 
 */
public class DynamicReplacementTable extends ReplacementTable
{

	
	/**
	 * The Replacers to apply
	 */
	private List<Replacer> replacers = new ArrayList<>();

	/**
	 * A {@link ReplacementTable} with support for {@link Replacer}s
	 * 
	 * @param table see {@link ReplacementTable#ReplacementTable(ITable, Map, Map, String, String)}
	 * @param objectMap see {@link ReplacementTable#ReplacementTable(ITable, Map, Map, String, String)}
	 * @param substringMap see {@link ReplacementTable#ReplacementTable(ITable, Map, Map, String, String)}
	 * @param startDelimiter see {@link ReplacementTable#ReplacementTable(ITable, Map, Map, String, String)}
	 * @param endDelimiter see {@link ReplacementTable#ReplacementTable(ITable, Map, Map, String, String)}
	 * @param replacers the replacers to be applied
	 */
	@SuppressWarnings("rawtypes")
	public DynamicReplacementTable(ITable table, Map objectMap, Map substringMap, String startDelimiter, String endDelimiter, List<Replacer> replacers)
	{
		super(table, objectMap, substringMap, startDelimiter, endDelimiter);
		this.replacers = replacers;
	}

	@Override
	public Object getValue(int row, String column) throws DataSetException
	{
		Object value = super.getValue(row, column);
		if (value instanceof String)
		{
			for(Replacer replacer : replacers){
				if(replacer.match((String)value)){
					return replacer.getReplacement();
				}
			}
		}
		return value;
	}
}

DynamicReplacmentDataSet.java

package io.github.tourniquet.replacer

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.dbunit.dataset.AbstractDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.ITableIterator;
import org.dbunit.dataset.ITableMetaData;
import org.dbunit.dataset.ReplacementDataSet;
import org.dbunit.dataset.ReplacementTable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * Mostly a copy of {@link ReplacementDataSet} to use {@link DynamicReplacementTable}
 *
 */
public class DynamicReplacementDataSet extends AbstractDataSet
{

	private ReplacementTable createReplacementTable(ITable table)
	{
		ReplacementTable replacementTable = new DynamicReplacementTable(table, _objectMap, _substringMap, _startDelim, _endDelim, replacers);
		replacementTable.setStrictReplacement(_strictReplacement);
		return replacementTable;
	}

	/**
	 * The replacers to be applied to the dataset
	 */
	private ArrayList<Replacer> replacers = new ArrayList<>();

	/**
	 * Add a replacer to be applied to the dataset
	 * 
	 * @param replacer the replacer to add
	 */
	public void addReplacer(Replacer replacer)
	{
		replacers.add(replacer);
	}

	public abstract static class Replacer
	{

		/**
		 * The pattern to match
		 */
		private Pattern pattern;

		/**
		 * The matcher
		 */
		private Matcher matcher;

		/**
		 * @param pattern the pattern to match
		 */
		public Replacer(Pattern pattern)
		{
			this.pattern = pattern;
		}

		/**
		 * @return the matcher
		 */
		protected Matcher getMatcher()
		{
			return matcher;
		}

		/**
		 * @param value the candidate for substitution
		 * @return whether the value matches the pattern
		 */
		public boolean match(String value)
		{
			matcher = pattern.matcher(value);
			return matcher.matches();
		}

		/**
		 * generates the replacement object for the previously matched value.
		 * 
		 * @return the substitution for the previously matched value
		 */
		public abstract Object getReplacement();
	}

	/*
	 * stuff below is just copied from original
	 * //////////////////////////////////////////////////////////////////////////
	 */

	private static final Logger logger = LoggerFactory.getLogger(DynamicReplacementDataSet.class);

	private final IDataSet _dataSet;

	private final Map _objectMap;

	private final Map _substringMap;

	private String _startDelim;

	private String _endDelim;

	private boolean _strictReplacement;

	/**
	 * Create a new ReplacementDataSet object that decorates the specified dataset.
	 * 
	 * @param dataSet the decorated table
	 */
	public DynamicReplacementDataSet(IDataSet dataSet)
	{
		_dataSet = dataSet;
		_objectMap = new HashMap();
		_substringMap = new HashMap();
	}

	/**
	 * Create a new ReplacementDataSet object that decorates the specified dataset.
	 * 
	 * @param dataSet the decorated dataset
	 * @param objectMap the replacement objects mapping
	 * @param substringMap the replacement substrings mapping
	 */
	public DynamicReplacementDataSet(IDataSet dataSet, Map objectMap, Map substringMap)
	{
		_dataSet = dataSet;
		_objectMap = objectMap == null ? new HashMap() : objectMap;
		_substringMap = substringMap == null ? new HashMap() : substringMap;
	}

	/**
	 * Setting this property to true indicates that when no replacement
	 * is found for a delimited substring the replacement will fail fast.
	 * 
	 * @param strictReplacement true if replacement should be strict
	 */
	public void setStrictReplacement(boolean strictReplacement)
	{
		this._strictReplacement = strictReplacement;
	}

	/**
	 * Add a new Object replacement mapping.
	 * 
	 * @param originalObject the object to replace
	 * @param replacementObject the replacement object
	 */
	public void addReplacementObject(Object originalObject, Object replacementObject)
	{
		logger.debug("addReplacementObject(originalObject={}, replacementObject={}) - start", originalObject, replacementObject);

		_objectMap.put(originalObject, replacementObject);
	}

	/**
	 * Add a new substring replacement mapping.
	 * 
	 * @param originalSubstring the substring to replace
	 * @param replacementSubstring the replacement substring
	 */
	public void addReplacementSubstring(String originalSubstring, String replacementSubstring)
	{
		logger.debug("addReplacementSubstring(originalSubstring={}, replacementSubstring={}) - start", originalSubstring, replacementSubstring);

		if (originalSubstring == null || replacementSubstring == null)
		{
			throw new NullPointerException();
		}

		_substringMap.put(originalSubstring, replacementSubstring);
	}

	/**
	 * Sets substring delimiters.
	 */
	public void setSubstringDelimiters(String startDelimiter, String endDelimiter)
	{
		logger.debug("setSubstringDelimiters(startDelimiter={}, endDelimiter={}) - start", startDelimiter, endDelimiter);

		if (startDelimiter == null || endDelimiter == null)
		{
			throw new NullPointerException();
		}

		_startDelim = startDelimiter;
		_endDelim = endDelimiter;
	}

	// //////////////////////////////////////////////////////////////////////////
	// AbstractDataSet class

	protected ITableIterator createIterator(boolean reversed) throws DataSetException
	{
		if (logger.isDebugEnabled())
			logger.debug("createIterator(reversed={}) - start", String.valueOf(reversed));

		return new ReplacementIterator(reversed ? _dataSet.reverseIterator() : _dataSet.iterator());
	}

	// //////////////////////////////////////////////////////////////////////////
	// IDataSet interface

	public String[] getTableNames() throws DataSetException
	{
		logger.debug("getTableNames() - start");

		return _dataSet.getTableNames();
	}

	public ITableMetaData getTableMetaData(String tableName) throws DataSetException
	{
		logger.debug("getTableMetaData(tableName={}) - start", tableName);

		return _dataSet.getTableMetaData(tableName);
	}

	public ITable getTable(String tableName) throws DataSetException
	{
		logger.debug("getTable(tableName={}) - start", tableName);

		return createReplacementTable(_dataSet.getTable(tableName));
	}

	// //////////////////////////////////////////////////////////////////////////
	// ReplacementIterator class

	private class ReplacementIterator implements ITableIterator
	{

		/**
		 * Logger for this class
		 */
		private final Logger logger = LoggerFactory.getLogger(ReplacementIterator.class);

		private final ITableIterator _iterator;

		public ReplacementIterator(ITableIterator iterator)
		{
			_iterator = iterator;
		}

		// //////////////////////////////////////////////////////////////////////
		// ITableIterator interface

		public boolean next() throws DataSetException
		{
			logger.debug("next() - start");

			return _iterator.next();
		}

		public ITableMetaData getTableMetaData() throws DataSetException
		{
			logger.debug("getTableMetaData() - start");

			return _iterator.getTableMetaData();
		}

		public ITable getTable() throws DataSetException
		{
			logger.debug("getTable() - start");

			return createReplacementTable(_iterator.getTable());
		}
	}
}

SampleUsage.java

/*
 * Example usage of DynamicReplacementDataSet
 */

DynamicReplacementDataSet dataSet;

private static final Pattern DATE_PATTERN = Pattern.compile("\\[(sysdate|systime|systimestamp)([\\+-](\\d+))?\\]");

private static final Pattern LASTOFMONTH_PATTERN = Pattern.compile("\\[(lastofmonth)([\\+-](\\d+))?\\]");

private static final Pattern SEQUENCE_PATTERN = Pattern.compile("\\[([\\w_]+_seq).nextval\\d*\\]");

public void setUpDatabase(String dataFile){
	InputStream dataStream = this.getClass().getClassLoader().getResourceAsStream(dataFile);
	if (dataStream != null) {
		FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
		builder.setColumnSensing(true);
		builder.setCaseSensitiveTableNames(true);
		dataSet = new RegappReplacementDataSet(builder.build(this.getClass().getClassLoader().getResourceAsStream(dataFile)));

		JndiDatabaseTester dbTester = new JndiDatabaseTester("java:jboss/datasources/myDS");
		IDatabaseConnection conn = dbTester.getConnection();
		initializeReplacements(conn);

		DatabaseOperation.REFRESH.execute(conn, dataSet);
	}

}

private void initializeReplacements(IDatabaseConnection conn) throws Exception {
	/*
	* Replace sysdate, systime, systimestamp
	* 
	* NOTE: [sysdate] uses the current system time,
	* while [sysdate+n] always sets hours, minutes and seconds to zero.
	*/
	dataSet.addReplacer(new Replacer(DATE_PATTERN) {
		@Override
		public Object getReplacement() {
			String offset_s = getMatcher().group(2);
			if (offset_s == null) {
				return new Date();
			} else {
				int offset = Integer.parseInt(offset_s);
				Calendar cal = getCalendarTodayZeroTime();
				cal.add(Calendar.DAY_OF_MONTH, offset);

				return cal.getTime();
			}
		}
	});

	/*
	* Replace lastofmonth
	* 
	* [lastofmonth] : last date of current month
	* [lastofmonth-2] : last date of two months earlier
	*/
	dataSet.addReplacer(new Replacer(LASTOFMONTH_PATTERN) {
		@Override
		public Object getReplacement() {
			String offset_s = getMatcher().group(2);
			int offset = 0;
			if (offset_s != null) {
				offset = Integer.parseInt(offset_s);
			} 

			Calendar cal = Calendar.getInstance();
			cal.add(Calendar.MONTH, offset);
			cal.set(Calendar.DAY_OF_MONTH, cal.getActualMaximum(Calendar.DAY_OF_MONTH));
			cal.set(Calendar.HOUR_OF_DAY, cal.getActualMaximum(Calendar.HOUR_OF_DAY));
			cal.set(Calendar.MINUTE, cal.getActualMaximum(Calendar.MINUTE));
			cal.set(Calendar.SECOND, cal.getActualMaximum(Calendar.SECOND));
			cal.set(Calendar.MILLISECOND, 0);

			return cal.getTime();
		}
	});

	//replace sequences
	final IDatabaseConnection replacerConn = conn;
	dataSet.addReplacer(new Replacer(SEQUENCE_PATTERN) {

		@Override
		public Object getReplacement() {

			Long value;
			String sequence = getMatcher().group(1);

			try {
				Statement statement = replacerConn.getConnection().createStatement();
				ResultSet result = statement.executeQuery("select "+sequence+".nextval from dual");
				result.next();
				value = new Long(result.getString(1));
			} catch (SQLException e) {
				if (e.getMessage().startsWith("ORA-02289: sequence does not exist")) {
					value = null;
				} else {
					throw new RuntimeException("Error running Replacer", e);
				}
		}
			return value;
		}
	});

	//odds and ends
	dataSet.addReplacementObject("[NULL]", null);
	dataSet.addReplacementObject("[null]", null);
}