Cacher is the code snippet organizer for pro developers

package jp.co.pioneer.iotcloud.lambda.rule.annotation;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;

@Retention(RetentionPolicy.RUNTIME)
public @interface TestExpectedDatabase {
	String value();
	TestDatabaseAssertionMode assertionMode() default TestDatabaseAssertionMode.DEFAULT;
}

package jp.co.pioneer.iotcloud.lambda.rule.annotation;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;

@Retention(RetentionPolicy.RUNTIME)
public @interface TestDatabaseSetup {
	String value();

	TestDatabaseOperation type() default TestDatabaseOperation.CLEAN_INSERT;
}

package jp.co.pioneer.iotcloud.lambda.rule.annotation;

import org.dbunit.operation.DatabaseOperation;

public enum TestDatabaseOperation {
	NONE(DatabaseOperation.NONE),
	UPDATE(DatabaseOperation.UPDATE),
	INSERT(DatabaseOperation.INSERT),
	REFRESH(DatabaseOperation.REFRESH),
	DELETE(DatabaseOperation.DELETE),
	DELETE_ALL(DatabaseOperation.DELETE_ALL),
	TRUNCATE_TABLE(DatabaseOperation.TRUNCATE_TABLE),
	CLEAN_INSERT(DatabaseOperation.CLEAN_INSERT);

	private DatabaseOperation databaseOperation;

	private TestDatabaseOperation(DatabaseOperation databaseOperation) {
		this.databaseOperation = databaseOperation;
	}

	public DatabaseOperation operation() {
		return databaseOperation;
	}
}

package jp.co.pioneer.iotcloud.lambda.rule.annotation;

import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;

@Retention(RetentionPolicy.RUNTIME)
public @interface TestDatabaseBackup {
	String[] value() default {};
	boolean setup() default false;
}

package jp.co.pioneer.iotcloud.lambda.rule.annotation;

public enum TestDatabaseAssertionMode {
	DEFAULT,
	NON_STRICT
}

package rule;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.net.URL;
import java.sql.SQLException;
import java.util.HashSet;
import java.util.Properties;
import java.util.Set;

import rule.annotation.TestDatabaseAssertionMode;
import rule.annotation.TestDatabaseBackup;
import rule.annotation.TestDatabaseSetup;
import rule.annotation.TestExpectedDatabase;

import org.dbunit.Assertion;
import org.dbunit.JdbcDatabaseTester;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.Column;
import org.dbunit.dataset.Columns;
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.csv.CsvDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.flywaydb.core.Flyway;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

public class DBUnitTester extends JdbcDatabaseTester implements TestRule {
	private Flyway flyway = new Flyway();

	public DBUnitTester(Properties properties) throws ClassNotFoundException {
		super(properties.getProperty("drivername"),
				properties.getProperty("url"),
				properties.getProperty("username"),
				properties.getProperty("password"));

		flyway.setDataSource(properties.getProperty("url"),
				properties.getProperty("username"),
				properties.getProperty("password"),
				new String[]{});
		flyway.setSqlMigrationSuffix("_ddl.sql");
	}

	protected void before() throws Exception {
	}

	protected void after() throws Exception {
	}

	protected IDataSet createDataSet(File file) throws DataSetException, IOException {
		ReplacementDataSet dataSet = new ReplacementDataSet(new CsvDataSet(file));
		dataSet.addReplacementObject("(NULL)", null);
		return dataSet;
	}

	@Override
	public Statement apply(Statement base, Description description) {
		return new Statement() {
			@Override
			public void evaluate() throws Throwable {
				setDataSet(null);
				initDatabase();

				before();

				final TestDatabaseSetup testDatabaseSetup = description.getAnnotation(TestDatabaseSetup.class);
				if (testDatabaseSetup != null) {
					URL dataFileUrl = description.getTestClass().getResource(testDatabaseSetup.value());
					IDataSet dataSet = createDataSet(new File(dataFileUrl.toURI()));
					setDataSet(dataSet);
				}

				File backupFile = makeBackupFile(description);
				if (testDatabaseSetup != null) {
					try {
						setSetUpOperation(testDatabaseSetup.type().operation());
						onSetup();
					} catch (Exception e) {
						if (backupFile != null) {
							backupFile.delete();
						}
						throw e;
					}
				}

				try {
					base.evaluate();

					assertion(description);
				} finally {
					try {
						after();
					} finally {
						if (backupFile != null) {
							try {
								restore(backupFile);
							} finally {
								backupFile.delete();
							}
						}

						onTearDown();
					}
				}
			}
		};
	}

	protected void initDatabase() {
		flyway.clean();
		flyway.migrate();
	}

	protected File makeBackupFile(Description description) throws Exception {
		final TestDatabaseBackup testDatabaseBackup = description.getAnnotation(TestDatabaseBackup.class);
		if (testDatabaseBackup == null) {
			return null;
		}

		IDatabaseConnection connection = null;

		Set<String> tableNames = new HashSet<>();
		try {
			connection = getConnection();
			QueryDataSet partialDataSet = new QueryDataSet(connection);
			IDataSet dataSet = getDataSet();
			if (testDatabaseBackup.setup() && dataSet != null) {
				ITableIterator dataSetIt = dataSet.iterator();
				while (dataSetIt.next()) {
					String tableName = dataSetIt.getTable().getTableMetaData().getTableName();
					if (tableNames.add(tableName)) {
						partialDataSet.addTable(tableName);
					}
				}
			}

			for (String tableName : testDatabaseBackup.value()) {
				if (tableNames.add(tableName)) {
					partialDataSet.addTable(tableName);
				}
			}

			File file = null;
			if (!tableNames.isEmpty()) {
				file = File.createTempFile("backup", ".xml");
				try (FileOutputStream output = new FileOutputStream(file)) {
					FlatXmlDataSet.write(partialDataSet, output);
				} catch (Exception e) {
					file.delete();
					throw e;
				}
			}

			return file;
		} finally {
			if (connection != null) {
				try {
					connection.close();
				} catch (SQLException e) {
				}
			}
		}

	}

	protected void restore(File file) throws Exception {
		IDatabaseConnection connection = null;
		try {
			connection = getConnection();

			IDataSet dataSet = new FlatXmlDataSetBuilder().build(file);
			DatabaseOperation.CLEAN_INSERT.execute(connection, dataSet);
		} finally {
			if (connection != null) {
				connection.close();
			}
		}
	}

	protected void assertion(Description description) throws SQLException, Exception {
		final TestExpectedDatabase testExpectedDatabase = description.getAnnotation(TestExpectedDatabase.class);
		if (testExpectedDatabase == null) {
			return;
		}

		URL resultUrl = description.getTestClass().getResource(testExpectedDatabase.value());
		IDataSet expectedDataSet = createDataSet(new File(resultUrl.toURI()));

		IDataSet actualDataSet = getConnection().createDataSet();
		ITableIterator expectedDataSetIt = expectedDataSet.iterator();
		while (expectedDataSetIt.next()) {
			ITable expectedTable = expectedDataSetIt.getTable();
			ITable actualTable = actualDataSet.getTable(expectedTable.getTableMetaData().getTableName());

			if (testExpectedDatabase.assertionMode() == TestDatabaseAssertionMode.DEFAULT) {
				Assertion.assertEquals(expectedTable, actualTable);
			} else if (testExpectedDatabase.assertionMode() == TestDatabaseAssertionMode.NON_STRICT) {
				ITableMetaData actualTableMetaData = actualTable.getTableMetaData();
				ITableMetaData expectedTableMetaData = expectedTable.getTableMetaData();

				Column[] notSpecifiedInExpected = Columns.getColumnDiff(expectedTableMetaData, actualTableMetaData).getActual();
				Set<String> ignoreCols = new HashSet<>();
				for (Column column : notSpecifiedInExpected) {
					ignoreCols.add(column.getColumnName());
				}

				Assertion.assertEqualsIgnoreCols(expectedTable, actualTable, ignoreCols.toArray(new String[ignoreCols.size()]));
			}
		}
	}
}