Setting Up Embedded Cassandra On Spring Project

When we first started using Cassandra, we immediately realized there would be a long period of development and prototyping until we reach the stable repository classes we would use. It was a new technology, client demands were changing constantly (in Cassandra you do query based modeling which means our repository classes changed as well) and we needed to build up confidence that we are doing right thing. We needed a tool which will give us fast feedback using production-like database setup without starting server, doing API triggers and checking what had landed in our database.  

CassandraUnit gave us exactly what we needed, which was a Java utility test tool providing us with an ability to test drive our repository and schema models according to the requirements. Basically, it starts Embedded Cassandra before test methods, has the ability to create structure and allows the developer to run CQL queries against production-like database without constantly starting the application, and without the need to have full API designed. We have decided to go with Datastax driver, so CassandraUnit enabled us to explore driver API in TDD manner. We receive the requirement, write the integration test using connected Session to Embedded Cassandra and start implementing until we see it turn green. 

Setup in Spring project and cleaning up the database between tests took quite an effort. We did this in a couple of steps. First, we created TestCassandraConfiguration which is plain Spring configuration class for active profile tests, and it creates the necessary beans and initializes Embedded Cassandra. The main thing here is to create a Session bean which we can use in other test methods to connect to our database.

@Bean
public Session session() throws Exception {
    if (session == null) {
        initialize();
    }

    if (sessionProxy == null) {
        sessionProxy = new SessionProxy(session);
    }

    return sessionProxy;
}

For now, we will ignore the SessionProxy thing, which will be explained later. If we do not have Session defined, initialize() method will start Embedded Cassandra, create a cluster and initialize table structure and connect to it. Here is the initialize method as well:    

private void initialize() throws Exception {
    LOGGER.info("Starting embedded cassandra server");
    EmbeddedCassandraServerHelper.startEmbeddedCassandra("another-cassandra.yaml");

    LOGGER.info("Connect to embedded db");
    cluster = Cluster.builder().addContactPoints(contact_points).withPort(port).build();
    session = cluster.connect();

    LOGGER.info("Initialize keyspace");
    final CQLDataLoader cqlDataLoader = new CQLDataLoader(session);
    cqlDataLoader.load(new ClassPathCQLDataSet(CQL, false, true, keyspace));
}

We created this configuration as DisposableBean implementation which enabled us to use the destroy() method for total cleanup. The complete example with all details of the configuration file is available on our Github account as part of our Spring Showcase Application which demonstrates usage of our Cassanadra Schema Migration Tool among other things.     

The second step is creating a CassandraTestExecutionListener which can be used on classes which need Embedded Cassandra. This listener is implementing AbstractTestExecutionListener and is responsible for cleanup between tests. It has afterTestMethod() which is doing cleanup.    

@Override
public void afterTestMethod(final TestContext testContext) throws Exception {
    LOGGER.debug("AfterTest: clean embedded cassandra");
    final Session session = (Session) TestApplicationContext.getBean("session");
    for (final String table : tables(session)) {
        session.execute(String.format("TRUNCATE %s.%s;", KEYSPACE, table));
    }
    super.afterTestMethod(testContext);
}

Nothing special here, for all tables we have in cluster metadata which we are fetching from session we perform truncate, so we can have a clean state before the next test. An important middle step in the communication between listener and configuration class is to create TestApplicationContext which will hold the beans defined in configuration and which can be used to access the Session bean.    

The last step was loading this listener on each test method, and that was it, we had Embedded Cassandra working and our Session connected to it injected in all repository classes instead of the production Session which would be created from production code and would be connected to production Cassandra cluster.    

@TestExecutionListeners({CassandraTestExecutionListener.class,
    DependencyInjectionTestExecutionListener.class})
public class SomeRepositoryTest {
    ...
}

Back to the SessionProxy thing. This was an another tweak we did. With Cassandra we used executeAsync() as much as we could; whenever we didn’t need response immediately, we used it. The problem was to test this. First, we took the Thread.sleep() road and added pauses all over our tests to be sure something would be written in order to read and verify it. That slowed our test suite a lot and was not a particular solution since there were no guarantees write will finish in sleep amount of seconds. We stubbed ResultSetFuture which was a result of executeAsync() method on Session and created a proxy on top of Session which would override executeAsync() to work as regular execute. As a result, we got synchronous executions even when our repository method used executeAsync (only in tests of course) which made our life easier while testing. We did not have to worry anymore if the results are written and how long should we pause the execution. 

To sum it all up. Cassandra is not new to us anymore, but we still like to have repository classes integration tested against production Cassandra like environment. Cassandra is known for its schema which evolves, modeling decisions must be revisited occasionally, so this gives us confidence that we can refactor and change. Also, performing tests this way has provided us a lot of insights into driver API much faster because we can now play with API and see the results right away. Last, but not the least, we are starting our backend applications much less often when we are developing database models, which significantly speeds up the development.