Skip to content

Instantly share code, notes, and snippets.

@ryu1
Last active September 13, 2022 08:33
Show Gist options
  • Save ryu1/04064a4831d699ee204cfc0fab22f062 to your computer and use it in GitHub Desktop.
Save ryu1/04064a4831d699ee204cfc0fab22f062 to your computer and use it in GitHub Desktop.
package test;
import org.dbunit.DefaultOperationListener;
import org.dbunit.IDatabaseTester;
import org.dbunit.IOperationListener;
import org.dbunit.JdbcDatabaseTester;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ReplacementDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.operation.DatabaseOperation;
import org.junit.After;
import org.junit.Before;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.util.ArrayList;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
/**
* DBUnitを使用したテストケースのスーパークラス.
*/
public class DBUnitTestBase {
@SuppressWarnings("unused")
private static final Logger log = LoggerFactory.getLogger(DBUnitTestBase.class);
@SuppressWarnings("all")
private static final String SELECT_ALL_TABLE_NAMES_QUERY = "SELECT \n" +
" c.relname \n" +
"FROM \n" +
" pg_class c JOIN \n" +
" pg_namespace n ON n.oid = c.relnamespace \n" +
"WHERE \n" +
" n.nspname = 'public' \n" +
" AND \n" +
" c.relkind = 'r' \n" +
"ORDER BY c.relname";
private static final String TRUNCATE_ALL_TABLES_QUERY
= "TRUNCATE TABLE %s RESTART IDENTITY CASCADE";
private final String pathOfResourceDirectory;
protected IDatabaseTester databaseTester;
/**
* コンスタラクタ.
* @param resourcePackage リソースファイルが配置されているパッケージ
*/
public DBUnitTestBase(final Package resourcePackage) {
super();
this.pathOfResourceDirectory = resourcePackage.getName().replace('.', '/');
}
private static IDatabaseTester createDatabaseTester(final String resource)
throws IOException, ClassNotFoundException {
Properties props;
try (InputStream propIs = Resources.getResourceAsStream(resource)) {
props = new Properties();
props.load(propIs);
}
JdbcDatabaseTester dbTester = new JdbcDatabaseTester(
"org.postgresql.Driver",
props.getProperty("url"),
props.getProperty("username"),
props.getProperty("password"), "public");
dbTester.setOperationListener(new CustomConfigurationOperationListener());
return dbTester;
}
/**
* Truncate all tables and reset all sequences.
*
* @param databaseTester databaseTester for target database
* @throws Exception SQLException etc.
*/
private static void resetDatabase(final IDatabaseTester databaseTester) throws Exception {
try (Connection conn = databaseTester.getConnection().getConnection()) {
// 該当テーブル取得
PreparedStatement statement = conn.prepareStatement(SELECT_ALL_TABLE_NAMES_QUERY);
statement.execute();
ResultSet resultSet = statement.getResultSet();
List<String> tables = new ArrayList<>();
while (resultSet.next()) {
tables.add(resultSet.getString(1));
}
statement.close();
final String truncate_query
= String.format(TRUNCATE_ALL_TABLES_QUERY, String.join(",", tables));
statement = conn.prepareStatement(truncate_query);
statement.executeUpdate();
statement.close();
}
}
@After
public void dbUnitTearDown() throws Exception {
if (databaseTester != null)
databaseTester.onTearDown();
}
private void setUpDataset(final IDatabaseTester databaseTester, final String filename)
throws Exception {
resetDatabase(databaseTester);
IDataSet set = createDataset(filename);
databaseTester.setDataSet(set);
databaseTester.setSetUpOperation(DatabaseOperation.CLEAN_INSERT);
databaseTester.onSetup();
}
/**
* 指定されたDatasetをデータベースに読み込みます.
* @param filename 読み込むデータセットのファイル
* @throws Exception database error, file not found etc
*/
@SuppressWarnings("SameParameterValue")
protected void setUpDataset(final String filename)
throws Exception {
databaseTester = createDatabaseTester("datasource.properties");
setUpDataset(databaseTester, filename);
}
/**
* データベースの全てのデータをdatasetにして返します.
* @return a data set
* @throws Exception database error
*/
@SuppressWarnings("SameParameterValue")
protected IDataSet createCurrentDatabaseDataset()
throws Exception {
IDatabaseConnection conn = null;
IDataSet databaseDataSet;
try {
conn = databaseTester.getConnection();
databaseDataSet = conn.createDataSet();
} catch (Exception e) {
if (conn != null)
conn.close();
throw e;
}
return databaseDataSet;
}
/**
* data setを作成します.
* @param filename Filename of XML-formatted data set file
* @return dataset
* @throws Exception database error
*/
private IDataSet buildDataset(final String filename) throws Exception {
final String resource = this.pathOfResourceDirectory + "/" + filename;
URL resourceUrl = this.getClass().getClassLoader().getResource(resource);
assert resourceUrl != null;
File resourceFile = new File(resourceUrl.toURI());
return new FlatXmlDataSetBuilder().build(resourceFile);
}
/**
* data setを作成します.
* @param filename Filename of XML-formatted data set file
* @return a data set
* @throws Exception database error, file not found etc
*/
protected IDataSet createDataset(final String filename)
throws Exception {
Map objectMap = new HashMap();
objectMap.put("[NULL]", null);
return new ReplacementDataSet(buildDataset(filename), objectMap, null);
}
/**
* data setを作成します.
* @param filename Filename of XML-formatted data set file
* @param objectMap Objects mapping to replace the dataset file
* @return a data set
* @throws Exception database error, file not found etc
*/
protected IDataSet createDataset(final String filename, final Map<Object, Object> objectMap)
throws Exception {
objectMap.put("[NULL]", null);
return new ReplacementDataSet(buildDataset(filename), objectMap, null);
}
public static class CustomConfigurationOperationListener extends DefaultOperationListener
implements IOperationListener {
@Override
public void connectionRetrieved(IDatabaseConnection iDatabaseConnection) {
super.connectionRetrieved(iDatabaseConnection);
iDatabaseConnection.getConfig().setProperty(DatabaseConfig.FEATURE_ALLOW_EMPTY_FIELDS, true);
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment