Dynamically Clearing Tables While Testing

I had to blog about this, or I’m going to forget about it. Isn’t that always the case? We write some utility, and run into some issues, and then 6 months from now we run into it again, but so many other things have come and gone that we’ve forgotten about them.

This has to do with integration testing, and clearing the state of the database between runs. In particular, clearing the contents of database tables. We’re not so much concerned with rebuilding the database — just the data that may be contained within it.

Generally, in our situation, we use liquibase to set up our database, and we use h2 for a fast, in-memory database in a Spring JPA codebase. Between tests, however, it can be a pain to always keep in mind what tables have data written to it, what foreign-key relationships need to be considered when deleting objects, etc. So why not have a utility class that just obtains the set of tables, ignores the FK constraints, and clears all the tables (except for liquibase).

Here’s what we came up with:

package com.sonicdrivein.drivein.services.util

import javax.persistence.EntityManager

import org.springframework.orm.jpa.JpaTransactionManager
import org.springframework.transaction.TransactionStatus
import org.springframework.transaction.support.DefaultTransactionDefinition

/**
  * Class containing test helper methods.
  */
 class DaoHelper {
 
     /**
      * Method to clear out all tables so that tests can be
      * run independently.
      */
     public static void clearTables(EntityManager em) {
 
         // Dynamically get all tables available to the dao and delete their contents, one by one.
         def schemas = getSchemas(em)
         schemas.each { schema ->
             def tables = getTables(em, schema)
             tables.each { table ->
                 if (table != "show tables") {
                     table = table[0]
 
                     System.out.print("Clearing data for table: ${schema}.${table} ... ")
                     // Don't muck with liquibase
                     if (table != 'DATABASECHANGELOG' && 
                         table != 'DATABASECHANGELOGLOCK') {
                         em.createNativeQuery("alter table ${schema}.${table} set referential_integrity false").executeUpdate()
                         em.createNativeQuery("delete from ${schema}.${table}").executeUpdate()
                         em.createNativeQuery("alter table ${schema}.${table} set referential_integrity true NOCHECK").executeUpdate()
                         System.out.println("done.")
                     } else {
                         System.out.println("skipped.")
                     }
                 }
             }
         }
     }
    /**
     * For the cases (rare) where we need to start/end transactions manually, this will definitely help.
     * You'll likely have to @Autowire a JpaTransactionManager into your test code, and pass it to this method.
     *
     * @param jtm The JpaTransactionManager
     * @return TransactionStatus
     */
     static TransactionStatus startTransaction( JpaTransactionManager jtm ) {
          DefaultTransactionDefinition transDef = new DefaultTransactionDefinition()

          return jtm.getTransaction(transDef)
     }

    /**
     * For the cases (rare) where we need to start/end transactions manually, this will definitely help.
     * Commit the transaction for this jtm & txstatus.
     *
     * @param jtm
     * @param txstatus
     */
     static void commitTransaction( JpaTransactionManager jtm, TransactionStatus txstatus ) {
          jtm.commit(txstatus)
     }

    /**
     * For the cases (rare) where we need to start/end transactions manually, this will definitely help.
     * In the context of this jtm, start a transaction, clear the tables, and then commit the transaction.
     *
     * @param jtm
     */
     static void clearTables( JpaTransactionManager jtm ) {
          TransactionStatus txstatus = startTransaction( jtm )
          EntityManager entityManager = jtm.entityManagerFactory.createEntityManager()
          entityManager.joinTransaction()

          clearTables(entityManager)

          commitTransaction(jtm, txstatus)
     }

    /**
     * For the given dao, return all existing schemas.
     * @param dao
     * @return
     */
     private static def getSchemas(EntityManager em) {
          def ret = em.createNativeQuery("show schemas").getResultList()
          def schemas = []
          ret.each { schema ->
            if (schema != "show schemas" &&
                schema != "INFORMATION_SCHEMA") {
                   schemas << schema 
            }
          }
          return schemas
     }

     /**
      * For the given dao, and schema, return all existing tables. 
      * @param dao 
      * @param schema 
      * @return 
      */ 
     private static def getTables(EntityManager em, String schema) { 
          def ret = em.createNativeQuery("show tables from ${schema}").getResultList() 
          def tables = [] 
          ret.each { table ->
               if (table != "show tables for ${schema}") {
                    tables << table
               }
          }

          return tables
    }
}

The use of this is quite simple. For most of our integration tests, we simply autowire in the EntityManager via @PersistenceContext:

@PersistenceContextEntityManager entityManager

Then, in the @After block, just invoke the DaoHelper to clear all your tables:

@Aftervoid tearDown() {     DaoHelper.clearTables( entityManager )}

And voila! The tables are cleared, and your tests are once again independent.

It’s a little more complicated if you’re not having Spring do all your transaction management; this is why the additional methods involvinga JpaTransactionManager are there. We had a situation where we are using both Mule, and Spring-JPA. Our integration test required that we spin up Mule… but that, in a testing context, caused some problems with the Spring integration. Our solution was to have our test case extend a base class based on this post by Ivan Krizsan:

package com.mycompany.services.util
 
 import org.mule.api.registry.RegistrationException
 import org.mule.api.registry.Registry
 import org.mule.tck.junit4.FunctionalTestCase
 import org.springframework.beans.factory.annotation.AutowiredAnnotationBeanPostProcessor
  
 // Kudos go to http://www.ivankrizsan.se/2015/12/13/spring-dependency-injection-in-mule-functionaltestcase/
 
 /**
  * Mule functional test case that will perform Spring dependency injection on the test class instance
  * prior to executing tests.
  * For dependency injection to be performed, the Mule context must be started before the {@code doSetUp}
  * method is executed.
  * For details on dependency injection performed, please refer to {@code AutowiredAnnotationBeanPostProcessor}.
  *
  * @author Ivan Krizsan
  * @see AutowiredAnnotationBeanPostProcessor
  */
 public abstract class SpringDIFunctionalTestCase extends FunctionalTestCase {
  
     @Override
     protected void doSetUp() throws RegistrationException {
         if (muleContext.isStarted()) {
             final Registry theRegistry = muleContext.getRegistry()
             final AutowiredAnnotationBeanPostProcessor theAutowiredAnnotationBeanPostProcessor
             try {
                 theAutowiredAnnotationBeanPostProcessor =
                     theRegistry.lookupObject(AutowiredAnnotationBeanPostProcessor.class)
             } catch (final RegistrationException theException) {
                 logger.error("An error occurred retrieving AutowiredAnnotationBeanPostProcessor", theException)
                 throw theException
             }
  
             if (theAutowiredAnnotationBeanPostProcessor != null) {
                 theAutowiredAnnotationBeanPostProcessor.processInjection(this)
             } else {
                 logger.warn("No AutowiredAnnotationBeanPostProcessor in the Mule registry, "
                     + "could not perform dependency injection on the test instance.")
             }
         } else {
             logger.warn("Mule context is not started, no dependency injection on test instance performed.")
         }
     }
 }

Using this, coupled with the DaoHelper, we can have our Mule integration test perform some autowiring … but transaction management within the test couldn’t reliably be done with @Transactional. Hence, the DaoHelper methods we added. Now our integration tests can extend SpringDIFunctionalTestCase, autowire the JpaTransactionManager, and we can invoke the DaoHelper.clearTables( jpaTransMgr ) method. We can also manually start and commit transactions quite easily.

Advertisements

About John Woodward

Principal Consultant at Improving Enterprises, Inc.
This entry was posted in Uncategorized. Bookmark the permalink.

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s