com.danhaywood.java.testsupport.dbunit.DbUnitRule.java Source code

Introduction

Here is the source code for com.danhaywood.java.testsupport.dbunit.DbUnitRule.java

Source

/*
 *  Copyright 2010-2013 Dan Haywood
 *
 *  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 com.danhaywood.java.testsupport.dbunit;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import java.nio.charset.Charset;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;

import org.dbunit.IDatabaseTester;
import org.dbunit.JdbcDatabaseTester;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.junit.rules.MethodRule;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.Statement;

import com.google.common.io.Resources;

public class DbUnitRule implements MethodRule {

    @Retention(RetentionPolicy.RUNTIME)
    @Target({ ElementType.METHOD })
    public static @interface Ddl {
        String[] value();
    }

    @Retention(RetentionPolicy.RUNTIME)
    @Target({ ElementType.METHOD })
    public static @interface JsonData {
        String value();
    }

    private final Class<?> resourceBase;

    private IDatabaseTester databaseTester;
    private IDatabaseConnection dbUnitConnection;

    private Connection connection;
    private java.sql.Statement statement;

    public DbUnitRule(Class<?> resourceBase, Class<?> driver, String url, String user, String password) {
        this.resourceBase = resourceBase;
        try {
            databaseTester = new JdbcDatabaseTester(driver.getName(), url, user, password);
            dbUnitConnection = databaseTester.getConnection();
            connection = dbUnitConnection.getConnection();
            statement = connection.createStatement();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    public Statement apply(final Statement base, final FrameworkMethod method, final Object target) {

        return new Statement() {

            @Override
            public void evaluate() throws Throwable {

                try {
                    Ddl ddl = method.getAnnotation(Ddl.class);
                    if (ddl != null) {
                        String[] values = ddl.value();
                        for (String value : values) {
                            executeUpdate(
                                    Resources.toString(resourceBase.getResource(value), Charset.defaultCharset()));
                        }
                    }

                    JsonData data = method.getAnnotation(JsonData.class);
                    if (data != null) {
                        IDataSet ds = new JsonDataSet(resourceBase.getResourceAsStream(data.value()));
                        databaseTester.setDataSet(ds);
                    }

                    databaseTester.onSetup();

                    base.evaluate();
                } finally {
                    databaseTester.onTearDown();
                }
            }
        };
    }

    public java.sql.Connection getConnection() {
        return connection;
    }

    public void executeUpdate(String sql) throws SQLException {
        statement.executeUpdate(sql);
    }

    public ResultSet executeQuery(String sql) throws SQLException {
        return statement.executeQuery(sql);
    }

    public IDataSet jsonDataSet(String datasetResource) {
        return new JsonDataSet(resourceBase.getResourceAsStream(datasetResource));
    }

    public ITable createQueryTable(String string, String string2) throws DataSetException, SQLException {
        return dbUnitConnection.createQueryTable(string, string2);
    }

}