DBUnitTestBase.java

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);
    }
  }

}