Session Mock Example

MockExample.java

import com.datastax.driver.core.*;
import org.mockito.Mockito;
import org.testng.annotations.Test;

import java.util.ArrayList;
import java.util.List;

/**
 * Created by atolbert on 1/26/15.
 */
public class MockExample {

    @Test
    public void testMockExample() throws Exception {
        Session session = Mockito.mock(Session.class);

        // Mock a ResultSet that gets returned from ResultSetFuture#get()
        ResultSet result = Mockito.mock(ResultSet.class);
        List<Row> rows = new ArrayList<Row>();
        Row r = Mockito.mock(Row.class);
        Mockito.doReturn(5).when(r).getInt(0);
        rows.add(r);

        Mockito.doReturn(rows).when(result).all();

        // Mock a ResultSetFuture that gets returned from Session#executeAsync()
        ResultSetFuture future = Mockito.mock(ResultSetFuture.class);
        Mockito.doReturn(result).when(future).get();
        Mockito.doReturn(future).when(session).executeAsync(Mockito.anyString());

        // Execute the query and print the 0th column of the first result.
        ResultSetFuture resultF = session.executeAsync("select value from table where key='a'");
        System.out.println(resultF.get().all().iterator().next().getInt(0));
    }
}