static ResultSetImpl helpTestBatching(StatementImpl statement, final int fetchSize, final int batchLength,
final int totalLength, final boolean partial) throws TeiidProcessingException, SQLException {
DQP dqp = statement.getDQP();
if (dqp == null) {
dqp = mock(DQP.class);
stub(statement.getDQP()).toReturn(dqp);
}
stub(statement.getFetchSize()).toReturn(fetchSize);
stub(dqp.processCursorRequest(Matchers.eq(REQUEST_ID), Matchers.anyInt(), Matchers.eq(fetchSize))).toAnswer(new Answer<ResultsFuture<ResultsMessage>>() {
@Override
public ResultsFuture<ResultsMessage> answer(
InvocationOnMock invocation) throws Throwable {
ResultsFuture<ResultsMessage> nextBatch = new ResultsFuture<ResultsMessage>();
int begin = Math.min(totalLength, (Integer)invocation.getArguments()[1]);
if (partial && begin == fetchSize + 1) {
begin = begin -5;
}
int length = Math.min(fetchSize, Math.min(totalLength - begin + 1, batchLength));
nextBatch.getResultsReceiver().receiveResults(exampleResultsMsg4(begin, length, begin + length - 1>= totalLength));
return nextBatch;
}
});
int initial = Math.min(fetchSize, batchLength);
ResultsMessage msg = exampleResultsMsg4(1, initial, initial == totalLength);
return new ResultSetImpl(msg, statement, new ResultSetMetaDataImpl(new MetadataProvider(DeferredMetadataProvider.loadPartialMetadata(msg.getColumnNames(), msg.getDataTypes())), null), 0);
}