Utilities for mocking DB calls

MockUtils.java

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;
  }
}