Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Utilities for mocking DB calls
import java.sql.Array;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
/**
* Helpers for mocking things up for tests.
*
* TODO move this into common ASAP.
*/
public class MockUtils {
/**
* Mocks result set a returning <tt>count</tt> rows.
*
* TODO Possibly it is better to add another version of {@link #mockRS(Map)} that takes a
* List<Map>, which will determine the result set.
*/
public static void setResultSetRowsToMock(ResultSet rs, int count) throws SQLException {
// In the answer closure below, only finals can be referenced
// but we want to mutate it. This way we'll have an {@link ResultSet#next()}
// return true the first time, and false next time, so we can get some results.
final List<Integer> countContainer = new ArrayList<>();
countContainer.add(count);
Mockito.when(rs.next()).thenAnswer(new Answer<Boolean>() {
@Override
public Boolean answer(InvocationOnMock invocation) throws Throwable {
int curCount = countContainer.get(0);
if (curCount < 0) {
return false;
}
countContainer.set(0, --curCount);
return true;
}
});
}
/**
* Mock an {@link java.sql.Array} to return the provided value.
*/
public static Array mockArrayReturningWhatever(Object[] whatever) throws SQLException {
Array retval = Mockito.mock(Array.class);
Mockito.when(retval.getArray()).thenReturn(whatever);
return retval;
}
/**
* Mock a ResultSet based on the provided Map of columns.
*
* TODO this can be refactored into some common helper... But would need to be in common, so a bit
* of a pain to do now.
*/
public static ResultSet mockRS(final Map<String, Object> map, final Map<Class, Object> defaults)
throws SQLException {
ResultSet rs = Mockito.mock(ResultSet.class);
// Now start mocking.
ArgumentCaptor<String> colNameArg = ArgumentCaptor.forClass(String.class);
// Potentially the below can all be combined into one piece of code with
// generics and reflection, but maybe later.
// Array
Mockito.when(rs.getArray(colNameArg.capture())).thenAnswer(new Answer<Array>() {
@Override
public Array answer(InvocationOnMock invocation) throws Throwable {
String colName = colNameArg.getValue();
if (map.containsKey(colName)) {
return (Array) map.get(colNameArg.getValue());
}
if (defaults != null && defaults.containsKey(Array.class)) {
return (Array) defaults.get(Array.class);
}
return null;
}
});
// Boolean
Mockito.when(rs.getBoolean(colNameArg.capture())).thenAnswer(new Answer<Boolean>() {
@Override
public Boolean answer(InvocationOnMock invocation) throws Throwable {
String colName = colNameArg.getValue();
if (map.containsKey(colName)) {
return (Boolean) map.get(colName);
}
if (defaults != null && defaults.containsKey(Boolean.class)) {
return (Boolean) defaults.get(Boolean.class);
}
return null;
}
});
// Double
Mockito.when(rs.getDouble(colNameArg.capture())).thenAnswer(new Answer<Double>() {
@Override
public Double answer(InvocationOnMock invocation) throws Throwable {
String colName = colNameArg.getValue();
if (map.containsKey(colName)) {
return (Double) map.get(colName);
}
if (defaults != null && defaults.containsKey(Double.class)) {
return (Double) defaults.get(Double.class);
}
return null;
}
});
// Integer
Mockito.when(rs.getInt(colNameArg.capture())).thenAnswer(new Answer<Integer>() {
@Override
public Integer answer(InvocationOnMock invocation) throws Throwable {
String colName = colNameArg.getValue();
if (map.containsKey(colName)) {
return (Integer) map.get(colName);
}
if (defaults != null && defaults.containsKey(Integer.class)) {
return (Integer) defaults.get(Integer.class);
}
return null;
}
});
// String
Mockito.when(rs.getString(colNameArg.capture())).thenAnswer(new Answer<String>() {
@Override
public String answer(InvocationOnMock invocation) throws Throwable {
String colName = colNameArg.getValue();
if (map.containsKey(colName)) {
return (String) map.get(colName);
}
if (defaults != null && defaults.containsKey(String.class)) {
return (String) defaults.get(String.class);
}
return null;
}
});
return rs;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment