Using Cassandra-unit with TestNG

I just started some tests with Cassandra and Spring Data Cassandra. I want to write some unit tests for this using TestNG. The Cassandra Unit project uses JUnit by default, but can be used with TestNG as well with some tweaking.

To start, add the following dependencies in your Maven pom:

<dependency>
            <groupId>org.springframework.data</groupId>
            <artifactId>spring-data-cassandra</artifactId>
            <version>1.0.0.RELEASE</version>
        </dependency>
        <dependency>
            <groupId>org.apache.cassandra</groupId>
            <artifactId>cassandra-all</artifactId>
            <version>2.0.5</version>
        </dependency>
        <dependency>
            <groupId>org.cassandraunit</groupId>
            <artifactId>cassandra-unit</artifactId>
            <version>2.0.2.1</version>
            <scope>test</scope>
            <exclusions>
                <exclusion>
                    <artifactId>cassandra-all</artifactId>
                    <groupId>org.apache.cassandra</groupId>
                </exclusion>
                <exclusion>
                    <groupId>com.datastax.cassandra</groupId>
                    <artifactId>cassandra-driver-core</artifactId>
                </exclusion>
                <exclusion>
                    <groupId>org.slf4j</groupId>
                    <artifactId>slf4j-log4j12</artifactId>
                </exclusion>
            </exclusions>
        </dependency>
        <dependency>
            <groupId>org.cassandraunit</groupId>
            <artifactId>cassandra-unit-spring</artifactId>
            <version>2.0.2.1</version>
            <scope>test</scope>
        </dependency>

Now we can setup Spring context loading with TestNG (using Java Config):

@Test
@ContextConfiguration
public class CassandraMessageRepositoryTest extends AbstractTestNGSpringContextTests
{

  public void testStoreMessage()
  {

  }

  @Configuration
  static class ContextConfiguration
  {
      // Add your @Bean's here
  }
}

After this, we need to start an embedded Cassandra before the Spring context loads, since the spring context needs to connect to the Cassandra server. For this, I have found no better way than overriding the @TestExecutionListeners. First create a new class:

import org.apache.cassandra.exceptions.ConfigurationException;
import org.apache.thrift.transport.TTransportException;
import org.cassandraunit.utils.EmbeddedCassandraServerHelper;
import org.springframework.beans.factory.config.AutowireCapableBeanFactory;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.DependencyInjectionTestExecutionListener;

import java.io.IOException;

public class PreContextLoadingTestExecutionListener extends DependencyInjectionTestExecutionListener
{
	protected void injectDependencies( final TestContext testContext ) throws Exception
	{
		doInitBeforeContextIsLoaded();

		Object bean = testContext.getTestInstance();
		AutowireCapableBeanFactory beanFactory = testContext.getApplicationContext().getAutowireCapableBeanFactory();
		beanFactory.autowireBeanProperties( bean, AutowireCapableBeanFactory.AUTOWIRE_NO, false );
		beanFactory.initializeBean( bean, testContext.getTestClass().getName() );
		testContext.removeAttribute( REINJECT_DEPENDENCIES_ATTRIBUTE );
	}

	private void doInitBeforeContextIsLoaded() throws InterruptedException, TTransportException, ConfigurationException, IOException
	{
		EmbeddedCassandraServerHelper.startEmbeddedCassandra();
	}
}

And register this listener in the test:

@Test
@ContextConfiguration
@TestExecutionListeners(inheritListeners = false, value = {DirtiesContextTestExecutionListener.class,
		PreContextLoadingTestExecutionListener.class})
public class CassandraMessageRepositoryTest extends AbstractTestNGSpringContextTests
{

}

The inner class ContextConfiguration looks like this in my example:

@Bean
		public MessageSourceService messageSourceService()
		{
			return mock( MessageSourceService.class );
		}

		@Bean
		public MessageRepository messageRepository() throws Exception
		{
			return new CassandraMessageRepository( cassandraTemplate(), messageSourceService(), objectMapper() );
		}

		@Bean
		public CassandraClusterFactoryBean cluster()
		{
			CassandraClusterFactoryBean cluster = new CassandraClusterFactoryBean();
			cluster.setContactPoints( "127.0.0.1" );
			cluster.setPort( 9142 );
			return cluster;
		}

		@Bean
		public CassandraSessionFactoryBean session() throws Exception
		{
			CassandraSessionFactoryBean session = new CassandraSessionFactoryBean();
			session.setCluster( cluster().getObject() );
			session.setConverter( converter() );
			session.setSchemaAction( SchemaAction.NONE );
			return session;
		}

		@Bean
		public CassandraOperations cassandraTemplate() throws Exception
		{
			return new CassandraTemplate( session().getObject() );
		}

		@Bean
		public CassandraMappingContext mappingContext()
		{
			return new BasicCassandraMappingContext();
		}

		@Bean
		public CassandraConverter converter()
		{
			return new MappingCassandraConverter( mappingContext() );
		}

		@Bean
		public ObjectMapper objectMapper()
		{
			ObjectMapper objectMapper = new ObjectMapper();
			objectMapper.registerModule( new MyCustomJacksonModule() );
			return objectMapper;
		}

A few points to note:
* CassandraMessageRepository is the class I want to test here
* The port of the embedded server is 9142, while the default port in a Cassandra installation is 9042.
* mock() is a Mockito static method to mock a collaborator to my class under test
* CassandraSessionFactoryBean should not declare the keyspace to use. A keyspace is created by cassandra-unit automatically.

With this in place, I can autowire what I need in my test:

	@Autowired
	private CassandraMessageRepository m_messageRepository;

	@Autowired
	private CassandraOperations m_cassandraOperations;

	@Autowired
	private Session m_session;

And then finally, the actual test:

	public void testStoreMessage()
	{
		ClassPathCQLDataSet dataSet = new ClassPathCQLDataSet( "com/mycompany/myapp/infrastructure/message/cassandra/create-event_message-table.cql" );
		CQLDataLoader loader = new CQLDataLoader( m_session );
		loader.load( dataSet );

		assertThat( m_cassandraOperations.queryForObject( "select count(*) from event_message", Long.class ) ).isEqualTo( 0 );

		m_messageRepository.storeMessage( Messages.createEventMessage() );

		assertThat( m_cassandraOperations.queryForObject( "select count(*) from event_message", Long.class ) ).isEqualTo( 1 );
	}

There are basically 4 steps in this method:
* First we load the CQL script from the classpath. For this we need the autowired Session
* We assert that the event_message table is empty (Using AssertJ statements)
* We store a message through our repository
* We check that there is 1 record in our table

If you have more then 1 test, you need to clean your database in between:

	@AfterMethod
	public void cleanCassandra() throws InterruptedException, TTransportException, ConfigurationException, IOException
	{
		EmbeddedCassandraServerHelper.cleanEmbeddedCassandra();
	}

It took me some time to figure all this out, so in case someone wants to use Cassandra Unit with TestNG, above is how to do it.

Advertisements

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