Using Cassandra unit with TestNG

Posted at — Jun 11, 2014

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.

Previous: Release JIRA Database Values plugin 3.3

Next: Opening multiple SSH sessions with iTerm...

Wim Deblauwe

Wim Deblauwe

Using Cassandra unit with TestNG