DB unit rule mysql
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()]));
}
}
}
}