Embedded Cassandra Unit test with bound statements

EmbeddedCassandraUnitTest.java

import com.datastax.driver.core.*;
import com.datastax.driver.core.querybuilder.QueryBuilder;
import org.apache.cassandra.exceptions.ConfigurationException;
import org.apache.thrift.transport.TTransportException;
import org.cassandraunit.utils.EmbeddedCassandraServerHelper;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.UUID;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;

public class EmbeddedCassandraUnitTest {

    private static final String KEYSPACE = "keystore";
    private static final String TABLE = "keys";
    private Session _cassandraSession;
    private BoundStatement _selectStatement;
    private BoundStatement _deleteStatement;
    private BoundStatement _insertStatement;

    @Before
    public void setup() throws ConfigurationException, TTransportException, IOException {
        EmbeddedCassandraServerHelper.startEmbeddedCassandra("/cassandra.yaml");
        connectCassandra();
    }
    
    private void connectCassandra() {
        Cluster cluster = Cluster.builder().addContactPoints("127.0.0.1").build();
        _cassandraSession = cluster.connect();
        this.createSchema();
        this.createStatements();
    }

    @Test
    public void testCRUD() throws Exception {
        assertNotNull(_cassandraSession.getCluster().getClusterName());

        // retrieve all rows to see if there is no data in cassandra
        ResultSet resultSet = _cassandraSession.execute(QueryBuilder.select().all().from(KEYSPACE, TABLE));
        assertEquals(0, resultSet.all().size());
        
        // create data for testing
        UUID id = UUID.randomUUID();
        UUID imageId = UUID.randomUUID();
        String keyString = "my-random-string";
        byte[] key = keyString.getBytes();
        ByteBuffer.wrap(key);

        // test insert, select, delete and select
        _cassandraSession.execute(_insertStatement.bind(id, imageId, ByteBuffer.wrap(key)));
        ResultSet sResultSet = _cassandraSession.execute(_selectStatement.bind(id, imageId));
        assertEquals(1, sResultSet.all().size());
        for (Row row : sResultSet.all()) {
            assertEquals(id, row.getUUID("id"));
            assertEquals(imageId, row.getUUID("imageId"));
            assertEquals(keyString, row.getBytes("key").toString());
        }
        _cassandraSession.execute(_deleteStatement.bind(id, imageId));
        ResultSet dResultSet = _cassandraSession.execute(_selectStatement.bind(id, imageId));
        assertEquals(0, dResultSet.all().size());

    }

    private void createStatements() {

        // create prepared statements
        PreparedStatement insertKey = _cassandraSession.prepare(
                QueryBuilder.insertInto("keystore", "keys")
                        .value("id", QueryBuilder.bindMarker())
                        .value("imageId", QueryBuilder.bindMarker())
                        .value("key", QueryBuilder.bindMarker()));

        PreparedStatement selectKey = _cassandraSession.prepare(QueryBuilder.select().all().from(KEYSPACE, TABLE)
                .where(QueryBuilder.eq("id", QueryBuilder.bindMarker()))
                .and(QueryBuilder.eq("imageId", QueryBuilder.bindMarker())));

        PreparedStatement deleteKey = _cassandraSession.prepare(QueryBuilder.delete().from(KEYSPACE, TABLE)
                .where(QueryBuilder.eq("id", QueryBuilder.bindMarker()))
                .and(QueryBuilder.eq("imageId", QueryBuilder.bindMarker())));

        // link prepared statements to boundstatements
        _insertStatement = new BoundStatement(insertKey);
        _selectStatement = new BoundStatement(selectKey);
        _deleteStatement = new BoundStatement(deleteKey);

    }

    private void createSchema() {
        // create keyspace if not exist
        _cassandraSession.execute("CREATE KEYSPACE IF NOT EXISTS " + KEYSPACE +
                " WITH replication={'class' : 'SimpleStrategy', 'replication_factor':1};");

        // create table if not exist
        _cassandraSession.execute("CREATE TABLE IF NOT EXISTS " + KEYSPACE + "." + TABLE + " (" +
                "id uuid," +
                "imageId uuid," +
                "key blob," +
                "PRIMARY KEY(id, imageId));");
    }

    @After
    public void tearDown() throws Exception {
        EmbeddedCassandraServerHelper.cleanEmbeddedCassandra();
    }
}