Java tutorial
/* * RHQ Management Platform * Copyright (C) 2005-2010 Red Hat, Inc. * All rights reserved. * * This program is free software; you can redistribute it and/or modify * it under the terms of the GNU General Public License as published by * the Free Software Foundation version 2 of the License. * * This program is distributed in the hope that it will be useful, * but WITHOUT ANY WARRANTY; without even the implied warranty of * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the * GNU General Public License for more details. * * You should have received a copy of the GNU General Public License * along with this program; if not, write to the Free Software * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA. */ package org.rhq.helpers.perftest.support.testng; import java.io.File; import java.io.FileNotFoundException; import java.io.IOException; import java.io.InputStream; import java.lang.annotation.Annotation; import java.lang.reflect.Method; import java.sql.Connection; import java.sql.SQLException; import java.sql.Statement; import java.util.Date; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.dbunit.database.DatabaseConfig; import org.dbunit.database.DatabaseDataSourceConnection; import org.dbunit.database.IDatabaseConnection; import org.dbunit.dataset.datatype.IDataTypeFactory; import org.rhq.core.db.ant.dbupgrade.SST_DirectSQL; import org.rhq.helpers.perftest.support.FileFormat; import org.rhq.helpers.perftest.support.Importer; import org.rhq.helpers.perftest.support.Input; import org.rhq.helpers.perftest.support.dbsetup.DbSetup; import org.rhq.helpers.perftest.support.input.FileInputStreamProvider; import org.rhq.helpers.perftest.support.input.InputStreamProvider; import org.testng.IInvokedMethod; import org.testng.IInvokedMethodListener; import org.testng.ITestResult; import org.testng.annotations.AfterMethod; import org.testng.annotations.AfterSuite; import org.testng.annotations.AfterTest; import org.testng.annotations.BeforeMethod; import org.testng.annotations.BeforeSuite; import org.testng.annotations.BeforeTest; import javax.naming.InitialContext; /** * An {@link IInvokedMethodListener method listener} that performs the database setup * for appropriately annotated test methods. * To add database setup support to a test class, annotate the class with * <code> * @Listeners({org.rhq.helpers.perftest.support.testng.DatabaseSetupInterceptor.class}) * </code> * * @author Lukas Krejci * @author Heiko W. Rupp */ public class DatabaseSetupInterceptor implements IInvokedMethodListener { private static final Log LOG = LogFactory.getLog(DatabaseSetupInterceptor.class); public void beforeInvocation(IInvokedMethod method, ITestResult testResult) { DatabaseState state = getRequiredDatabaseState(method); if (state == null) { return; } String dbUrl = "-unknown-"; Connection jdbcConnection = null; Statement statement = null; try { InputStreamProvider streamProvider = getInputStreamProvider(state.url(), state.storage(), method); IDatabaseConnection connection = new DatabaseDataSourceConnection(new InitialContext(), "java:/RHQDS"); jdbcConnection = connection.getConnection(); dbUrl = jdbcConnection.getMetaData().getURL(); System.out.println("Using database at " + dbUrl); System.out.flush(); setDatabaseType(connection); try { statement = jdbcConnection.createStatement(); statement.execute("DROP TABLE RHQ_SUBJECT CASCADE"); } catch (SQLException e) { System.out.println("Don't worry about : " + e.getMessage()); } finally { if (statement != null) statement.close(); } try { statement = jdbcConnection.createStatement(); statement.execute("DROP TABLE RHQ_CONFIG CASCADE"); } catch (SQLException e) { System.out.println("Don't worry about : " + e.getMessage()); } finally { if (statement != null) statement.close(); } System.out.flush(); FileFormat format = state.format(); Input input = format.getInput(streamProvider); try { DbSetup dbSetup = new DbSetup(jdbcConnection); dbSetup.setup(state.dbVersion()); Importer.run(connection, input); dbSetup.upgrade(null); } finally { input.close(); } } catch (Exception e) { LOG.warn("Failed to setup a database at [ " + dbUrl + "] for method '" + method.getTestMethod().getMethodName() + "'.", e); } finally { if (statement != null) try { statement.close(); } catch (SQLException e) { LOG.error("Failed to close a statement: " + e.getMessage()); } if (jdbcConnection != null) try { jdbcConnection.close(); } catch (SQLException e) { LOG.error("Failed to close a JDBC connetion: " + e.getMessage()); } } System.out.flush(); System.err.flush(); } private void setDatabaseType(IDatabaseConnection connection) throws SQLException { DatabaseConfig config = connection.getConfig(); String name = connection.getConnection().getMetaData().getDatabaseProductName().toLowerCase(); int major = connection.getConnection().getMetaData().getDatabaseMajorVersion(); IDataTypeFactory type = null; if (name.contains("postgres")) { type = new org.dbunit.ext.postgresql.PostgresqlDataTypeFactory(); } else if (name.contains("oracle")) { if (major >= 10) { type = new org.dbunit.ext.oracle.Oracle10DataTypeFactory(); } else { type = new org.dbunit.ext.oracle.OracleDataTypeFactory(); } } if (type != null) { LOG.info("setting db type for dbunit to " + type.getClass().getCanonicalName()); config.setProperty("http://www.dbunit.org/properties/datatypeFactory", type); } } /* Dropping the tables explicitly is a workaround around a problem in the perf tests * but otherwise it should not be necessary at all. * Even then, dropping the table in the beforeInvocation() should be enough and we * shouldn't have to drop it again in the afterInvocation(). That is going to leave * the database corrupted because it is going to miss this table altogether. */ public void afterInvocation(IInvokedMethod method, ITestResult testResult) { DatabaseState state = getRequiredDatabaseState(method); if (state == null) { return; } Connection jdbcConnection = null; Statement statement = null; try { IDatabaseConnection connection = new DatabaseDataSourceConnection(new InitialContext(), "java:/RHQDS"); jdbcConnection = connection.getConnection(); statement = jdbcConnection.createStatement(); statement.execute("DROP TABLE RHQ_SUBJECT CASCADE"); } catch (Exception e) { System.err.println("== drop subject table failed: " + e.getMessage()); } finally { if (statement != null) try { statement.close(); } catch (SQLException e) { LOG.error("Failed to close a statement: " + e.getMessage()); } if (jdbcConnection != null) try { jdbcConnection.close(); } catch (SQLException e) { LOG.error("Failed to close a JDBC connection: " + e.getMessage()); } } } /** * Obtain the required database state by looking for the @DatabaseState annotation. * Lookup is first done at method level and if not found done at class level. * @param method Method that TestNG is about to invoke * @return the desired database state or null if @DatabaseState is not given. * @see org.rhq.helpers.perftest.support.testng.DatabaseState */ private static DatabaseState getRequiredDatabaseState(IInvokedMethod method) { Method javaMethod = method.getTestMethod().getMethod(); DatabaseState annotation = javaMethod.getAnnotation(DatabaseState.class); if (annotation == null) { // System.out.println("Method : " + javaMethod.getName()); boolean skip = false; // Filter out methods that are marked as setup/tear down Annotation[] annots = javaMethod.getAnnotations(); for (Annotation an : annots) { // System.out.println(" : " + an.toString()); if (an.annotationType().equals(BeforeMethod.class) || an.annotationType().equals(AfterMethod.class) || an.annotationType().equals(BeforeSuite.class) || an.annotationType().equals(AfterSuite.class) || an.annotationType().equals(BeforeTest.class) || an.annotationType().equals(AfterTest.class)) skip = true; } if (!skip) annotation = javaMethod.getDeclaringClass().getAnnotation(DatabaseState.class); // else // System.out.println(" ..... Skipped"); } return annotation; } private static InputStreamProvider getInputStreamProvider(final String url, DatabaseStateStorage storage, final IInvokedMethod method) throws FileNotFoundException { switch (storage) { case CLASSLOADER: return new InputStreamProvider() { public InputStream createInputStream() throws IOException { ClassLoader cl = method.getTestMethod().getMethod().getDeclaringClass().getClassLoader(); return cl.getResourceAsStream(url); } }; case FILESYSTEM: return new FileInputStreamProvider(new File(url)); default: return null; } } }