/sigmah/src/test/java/org/sigmah/server/dao/hibernate/LoadDataSet.java

  1/*
  2 * All Sigmah code is released under the GNU General Public License v3
  3 * See COPYRIGHT.txt and LICENSE.txt.
  4 */
  5
  6package org.sigmah.server.dao.hibernate;
  7
  8import org.dbunit.DatabaseUnitException;
  9import org.dbunit.database.DatabaseConnection;
 10import org.dbunit.database.IDatabaseConnection;
 11import org.dbunit.dataset.DataSetException;
 12import org.dbunit.dataset.IDataSet;
 13import org.dbunit.dataset.xml.FlatXmlDataSet;
 14import org.dbunit.ext.h2.H2Connection;
 15import org.dbunit.ext.hsqldb.HsqldbConnection;
 16import org.dbunit.ext.mssql.InsertIdentityOperation;
 17import org.dbunit.ext.mssql.MsSqlConnection;
 18import org.dbunit.operation.DatabaseOperation;
 19import org.hibernate.Session;
 20import org.hibernate.ejb.HibernateEntityManager;
 21import org.hibernate.jdbc.Work;
 22import org.junit.internal.runners.model.MultipleFailureException;
 23import org.junit.runners.model.Statement;
 24import org.sigmah.server.dao.DatabaseCleaner;
 25
 26import javax.persistence.EntityManager;
 27import javax.persistence.EntityManagerFactory;
 28import java.io.IOException;
 29import java.io.InputStream;
 30import java.io.InputStreamReader;
 31import java.sql.Connection;
 32import java.sql.SQLException;
 33import java.util.ArrayList;
 34import java.util.List;
 35
 36public class LoadDataSet extends Statement {
 37    private final Statement next;
 38    private final Object target;
 39    private final EntityManagerFactory emf;
 40    private String name;
 41
 42
 43    public LoadDataSet(EntityManagerFactory emf, Statement next, String name, Object target) {
 44        this.next = next;
 45        this.target = target;
 46        this.emf = emf;
 47        this.name = name;
 48    }
 49
 50    @Override
 51    public void evaluate() throws Throwable {
 52        System.err.println("DBUnit: removing all rows");
 53        removeAllRows();
 54
 55        System.err.println("DBUnit: loading " + name + " into the database.");
 56        IDataSet data = loadDataSet();
 57
 58        List<Throwable> errors = new ArrayList<Throwable>();
 59        errors.clear();
 60        try {
 61            populate(data);
 62            next.evaluate();
 63        } catch (Throwable e) {
 64            errors.add(e);
 65        }
 66        MultipleFailureException.assertEmpty(errors);
 67    }
 68
 69    private IDataSet loadDataSet() throws IOException, DataSetException {
 70        InputStream in = target.getClass().getResourceAsStream(name);
 71        if (in == null) {
 72            throw new Error("Could not find resource '" + name + "'");
 73        }
 74        return new FlatXmlDataSet(new InputStreamReader(in), false, true, false);
 75    }
 76
 77    private void populate(final IDataSet dataSet) {
 78        executeOperation(InsertIdentityOperation.INSERT, dataSet);
 79    }
 80
 81    private void removeAllRows() {
 82        EntityManager em = emf.createEntityManager();
 83        DatabaseCleaner cleaner = new DatabaseCleaner(em, SQLDialectProvider.from(em));
 84        cleaner.clean();
 85        em.close();
 86    }
 87
 88    private void executeOperation(final DatabaseOperation op, final IDataSet dataSet) {
 89        HibernateEntityManager hem = (HibernateEntityManager) emf.createEntityManager();
 90        Session session = hem.getSession();
 91        session.doWork(new Work() {
 92            @Override
 93            public void execute(Connection connection) throws SQLException {
 94                try {
 95                    IDatabaseConnection dbUnitConnection = createDbUnitConnection(connection);
 96                    op.execute(dbUnitConnection, dataSet);
 97                } catch (DatabaseUnitException e) {
 98                    throw new RuntimeException(e);
 99                }
100            }
101        });
102        hem.close();
103    }
104
105    private IDatabaseConnection createDbUnitConnection(Connection connection) throws DatabaseUnitException, SQLException {
106        String dbName = connection.getMetaData().getDatabaseProductName();
107        if(dbName.equals("Microsoft SQL Server")) {
108            IDatabaseConnection con = new MsSqlConnection(connection);
109            con.getConfig().setProperty("http://www.dbunit.org/properties/escapePattern", "[?]");
110            return con;
111        } else if(dbName.equals("HSQL Database Engine")) {
112            return new HsqldbConnection(connection, null);
113        } else if(dbName.equals("H2")) {
114            return new H2Connection(connection, null);
115        } else {
116            return new DatabaseConnection(connection);
117        }
118    }
119}