[utils] / system / spring / src / test / java / org / wamblee / system / spring / component / DatabaseTesterComponent.java

/*
 * Copyright 2005-2010 the original author or authors.
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */ 
package org.wamblee.system.spring.component;

import junit.framework.TestCase;

import org.dbunit.DatabaseUnitException;

import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.DatabaseSequenceFilter;
import org.dbunit.database.IDatabaseConnection;

import org.dbunit.dataset.FilteredDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.filter.ITableFilter;

import org.dbunit.operation.DatabaseOperation;

import org.hibernate.SessionFactory;

import org.springframework.beans.factory.NoSuchBeanDefinitionException;

import org.springframework.context.ApplicationContext;

import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.datasource.DataSourceUtils;

import org.springframework.orm.hibernate3.HibernateTemplate;

import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionDefinition;
import org.springframework.transaction.support.TransactionCallback;
import org.springframework.transaction.support.TransactionCallbackWithoutResult;
import org.springframework.transaction.support.TransactionTemplate;

import org.wamblee.test.spring.TestTransactionCallback;
import org.wamblee.test.spring.TestTransactionCallbackWithoutResult;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.logging.Logger;

import javax.sql.DataSource;

/**
 * Test support class for database testing. Currently, this still requires the
 * spring platform transaction manager and hibernate template.
 */
public class DatabaseTesterComponent {
    private static final Logger LOG = Logger
        .getLogger(DatabaseTesterComponent.class.getName());

    /**
     * Schema pattern.
     */
    private static final String SCHEMA_PATTERN = "%";

    /**
     * Cached spring application context.
     */
    private ApplicationContext context;

    private HibernateTemplate hibernateTemplate;

    private PlatformTransactionManager transactionManager;

    private DataSource dataSource;

    /**
     * Creates a new DatabaseTesterComponent object.
     * 
     */
    public DatabaseTesterComponent(HibernateTemplate aHibernateTemplate,
        PlatformTransactionManager aTransactionManager, DataSource aDataSource) {
        hibernateTemplate = aHibernateTemplate;
        transactionManager = aTransactionManager;
        dataSource = aDataSource;
    }

    /**
     * 
     * @return Hibernate session factory.
     */
    protected SessionFactory getSessionFactory() {
        return hibernateTemplate.getSessionFactory();
    }

    /**
     * Performs common initialization for test cases:
     * <ul>
     * <li>Cleaning the database.</li>
     * </ul>
     * 
     * @throws Exception
     */
    public void setUp() throws Exception {
        LOG.info("Performing setUp()");

        cleanDatabase();
    }

    /**
     * Performs common tear down after execution of a test case. Currenlty this
     * method does nothing.
     * 
     * @throws Exception
     */
    protected void tearDown() throws Exception {
        // Empty
    }

    /**
     * 
     * @return Transaction manager
     */
    protected PlatformTransactionManager getTransactionManager() {
        return transactionManager;
    }

    /**
     * 
     * @return Starts a new transaction.
     */
    protected TransactionStatus getTransaction() {
        DefaultTransactionDefinition def = new DefaultTransactionDefinition();
        def.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRED);

        return getTransactionManager().getTransaction(def);
    }

    /**
     * Returns the hibernate template for executing hibernate-specific
     * functionality.
     * 
     * @return Hibernate template.
     */
    public HibernateTemplate getTemplate() {
        return hibernateTemplate;
    }

    /**
     * Flushes the session. Should be called after some Hibernate work and
     * before JDBC is used to check results.
     */
    public void flush() {
        getTemplate().flush();
    }

    /**
     * Flushes the session first and then removes all objects from the Session
     * cache. Should be called after some Hibernate work and before JDBC is used
     * to check results.
     */
    public void clear() {
        flush();
        getTemplate().clear();
    }

    /**
     * Evicts the object from the session. This is essential for the
     * implementation of unit tests where first an object is saved and is
     * retrieved later. By removing the object from the session, Hibernate must
     * retrieve the object again from the database.
     * 
     * @param aObject
     */
    protected void evict(Object aObject) {
        getTemplate().evict(aObject);
    }

    /**
     * Gets the connection.
     * 
     * @return Connection.
     */
    public Connection getConnection() {
        return DataSourceUtils.getConnection(getDataSource());
    }

    public void cleanDatabase() throws SQLException {
        if (!isDatabaseConfigured()) {
            return;
        }

        String[] tables = getTableNames();

        try {
            IDatabaseConnection connection = new DatabaseConnection(
                getConnection());
            ITableFilter filter = new DatabaseSequenceFilter(connection, tables);
            IDataSet dataset = new FilteredDataSet(filter, connection
                .createDataSet(tables));

            DatabaseOperation.DELETE_ALL.execute(connection, dataset);
        } catch (DatabaseUnitException e) {
            SQLException exc = new SQLException(e.getMessage());
            exc.initCause(e);
            throw exc;
        }
    }

    /**
     * 
     * 
     * @throws SQLException
     */
    public String[] getTableNames() throws SQLException {
        List<String> result = new ArrayList<String>();
        LOG.fine("Getting database table names to clean (schema: '" +
            SCHEMA_PATTERN + "'");

        ResultSet tables = getConnection().getMetaData().getTables(null,
            SCHEMA_PATTERN, "%", new String[] { "TABLE" });

        while (tables.next()) {
            String table = tables.getString("TABLE_NAME");

            // Make sure we do not touch hibernate's specific
            // infrastructure tables.
            if (!table.toLowerCase().startsWith("hibernate")) {
                result.add(table);
                LOG.fine("Adding " + table +
                    " to list of tables to be cleaned.");
            }
        }

        return (String[]) result.toArray(new String[0]);
    }

    /**
     * 
     * 
     * @throws SQLException
     */
    public void emptyTables(List aTableList) throws SQLException {
        Iterator liTable = aTableList.iterator();

        while (liTable.hasNext()) {
            emptyTable((String) liTable.next());
        }
    }

    /**
     * 
     * 
     * @throws SQLException
     */
    public void emptyTable(String aTable) throws SQLException {
        executeSql("delete from " + aTable);
    }

    /**
     * 
     * 
     * @throws SQLException
     */
    public void dropTable(String aTable) throws SQLException {
        executeQuery("drop table " + aTable);
    }

    /**
     * Executes an SQL statement within a transaction.
     * 
     * @param aSql
     *            SQL statement.
     * 
     * @return Return code of the corresponding JDBC call.
     */
    public int executeSql(final String aSql) {
        return executeSql(aSql, new Object[0]);
    }

    /**
     * Executes an SQL statement within a transaction. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            SQL statement.
     * @param aArg
     *            Argument of the sql statement.
     * 
     * @return Return code of the corresponding JDBC call.
     */
    public int executeSql(final String aSql, final Object aArg) {
        return executeSql(aSql, new Object[] { aArg });
    }

    /**
     * Executes an sql statement. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            SQL query to execute.
     * @param aArgs
     *            Arguments.
     * 
     * @return Number of rows updated.
     */
    public int executeSql(final String aSql, final Object[] aArgs) {
        Map results = executeTransaction(new TestTransactionCallback() {
            public Map execute() throws Exception {
                JdbcTemplate template = new JdbcTemplate(getDataSource());
                int result = template.update(aSql, aArgs);

                Map<String, Integer> map = new TreeMap<String, Integer>();
                map.put("result", new Integer(result));

                return map;
            }
        });

        return ((Integer) results.get("result")).intValue();
    }

    /**
     * Executes a transaction with a result.
     * 
     * @param aCallback
     *            Callback to do your transactional work.
     * 
     * @return Result.
     */
    public Object executeTransaction(TransactionCallback aCallback) {
        TransactionTemplate lTemplate = new TransactionTemplate(
            getTransactionManager());

        return lTemplate.execute(aCallback);
    }

    /**
     * Executes a transaction without a result.
     * 
     * @param aCallback
     *            Callback to do your transactional work. .
     */
    public void executeTransaction(TransactionCallbackWithoutResult aCallback) {
        TransactionTemplate template = new TransactionTemplate(
            getTransactionManager());
        template.execute(aCallback);
    }

    /**
     * Executes a transaction with a result, causing the testcase to fail if any
     * type of exception is thrown.
     * 
     * @param aCallback
     *            Code to be executed within the transaction.
     * 
     * @return Result.
     * 
     */
    public Map executeTransaction(final TestTransactionCallback aCallback) {
        return (Map) executeTransaction(new TransactionCallback() {
            public Object doInTransaction(TransactionStatus aArg) {
                try {
                    return aCallback.execute();
                } catch (Exception e) {
                    // test case must fail.
                    e.printStackTrace();
                    throw new RuntimeException(e);
                }
            }
        });
    }

    /**
     * Executes a transaction with a result, causing the testcase to fail if any
     * type of exception is thrown.
     * 
     * @param aCallback
     *            Code to be executed within the transaction.
     * 
     */
    public void executeTransaction(
        final TestTransactionCallbackWithoutResult aCallback) {
        executeTransaction(new TransactionCallbackWithoutResult() {
            public void doInTransactionWithoutResult(TransactionStatus aArg) {
                try {
                    aCallback.execute();
                } catch (Exception e) {
                    // test case must fail.
                    throw new RuntimeException(e.getMessage(), e);
                }
            }
        });
    }

    /**
     * Executes an SQL query.
     * 
     * @param aSql
     *            Query to execute.
     * 
     * @return Result set.
     */
    public ResultSet executeQuery(String aSql) {
        return executeQuery(aSql, new Object[0]);
    }

    /**
     * Executes a query with a single argument. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            Query.
     * @param aArg
     *            Argument.
     * 
     * @return Result set.
     */
    public ResultSet executeQuery(String aSql, Object aArg) {
        return executeQuery(aSql, new Object[] { aArg });
    }

    /**
     * Executes a query. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aSql
     *            Sql query.
     * @param aArgs
     *            Arguments to the query.
     * 
     * @return Result set.
     * 
     */
    public ResultSet executeQuery(final String aSql, final Object[] aArgs) {
        try {
            Connection connection = getConnection();

            PreparedStatement statement = connection.prepareStatement(aSql);
            setPreparedParams(aArgs, statement);

            return statement.executeQuery();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Sets the values of a prepared statement. See
     * {@link #setPreparedParam(int, PreparedStatement, Object)}for details on
     * supported argument types.
     * 
     * @param aArgs
     *            Arguments to the prepared statement.
     * @param aStatement
     *            Prepared statement
     * 
     * @throws SQLException
     */
    private void setPreparedParams(final Object[] aArgs,
        PreparedStatement aStatement) throws SQLException {
        for (int i = 1; i <= aArgs.length; i++) {
            setPreparedParam(i, aStatement, aArgs[i - 1]);
        }
    }

    /**
     * Sets a prepared statement parameter.
     * 
     * @param aIndex
     *            Index of the parameter.
     * @param aStatement
     *            Prepared statement.
     * @param aObject
     *            Value Must be of type Integer, Long, or String.
     * 
     * @throws SQLException
     */
    private void setPreparedParam(int aIndex, PreparedStatement aStatement,
        Object aObject) throws SQLException {
        if (aObject instanceof Integer) {
            aStatement.setInt(aIndex, ((Integer) aObject).intValue());
        } else if (aObject instanceof Long) {
            aStatement.setLong(aIndex, ((Integer) aObject).longValue());
        } else if (aObject instanceof String) {
            aStatement.setString(aIndex, (String) aObject);
        } else {
            TestCase.fail("Unsupported object type for prepared statement: " +
                aObject.getClass() + " value: " + aObject + " statement: " +
                aStatement);
        }
    }

    private boolean isDatabaseConfigured() {
        try {
            getDataSource();
        } catch (NoSuchBeanDefinitionException e) {
            return false;
        }

        return true;
    }

    /**
     * 
     * @return Returns the dataSource.
     */
    public DataSource getDataSource() {
        return dataSource;
    }

    /**
     * 
     * 
     * @return
     * 
     * @throws SQLException
     */
    public int getTableSize(final String aTable) throws SQLException {
        ResultSet resultSet = executeQuery("select * from " + aTable);
        int count = 0;

        while (resultSet.next()) {
            count++;
        }

        return count;
    }

    public int countResultSet(ResultSet aResultSet) throws SQLException {
        int count = 0;

        while (aResultSet.next()) {
            count++;
        }

        return count;
    }
}