final Query query,
Object... args // does this assume query intervals must be ordered?
)
{
if (args.length % 2 != 0) {
throw new ISE("args.length must be divisible by two, was %d", args.length);
}
final List<Interval> queryIntervals = Lists.newArrayListWithCapacity(args.length / 2);
final List<List<Iterable<Result<Object>>>> expectedResults = Lists.newArrayListWithCapacity(queryIntervals.size());
for (int i = 0; i < args.length; i += 2) {
final Interval interval = (Interval) args[i];
final Iterable<Result<Object>> results = (Iterable<Result<Object>>) args[i + 1];
if (queryIntervals.size() > 0 && interval.equals(queryIntervals.get(queryIntervals.size() - 1))) {
expectedResults.get(expectedResults.size() - 1).add(results);
} else {
queryIntervals.add(interval);
expectedResults.add(Lists.<Iterable<Result<Object>>>newArrayList(results));
}
}
for (int i = 0; i < queryIntervals.size(); ++i) {
List<Object> mocks = Lists.newArrayList();
mocks.add(serverView);
final Interval actualQueryInterval = new Interval(
queryIntervals.get(0).getStart(), queryIntervals.get(i).getEnd()
);
final List<Map<DruidServer, ServerExpectations>> serverExpectationList = populateTimeline(
queryIntervals,
expectedResults,
i,
mocks
);
List<Capture> queryCaptures = Lists.newArrayList();
final Map<DruidServer, ServerExpectations> finalExpectation = serverExpectationList.get(
serverExpectationList.size() - 1
);
for (Map.Entry<DruidServer, ServerExpectations> entry : finalExpectation.entrySet()) {
DruidServer server = entry.getKey();
ServerExpectations expectations = entry.getValue();
EasyMock.expect(serverView.getQueryRunner(server))
.andReturn(expectations.getQueryRunner())
.once();
final Capture<? extends Query> capture = new Capture();
final Capture<? extends Map> context = new Capture();
queryCaptures.add(capture);
QueryRunner queryable = expectations.getQueryRunner();
if (query instanceof TimeseriesQuery) {
List<String> segmentIds = Lists.newArrayList();
List<Interval> intervals = Lists.newArrayList();
List<Iterable<Result<TimeseriesResultValue>>> results = Lists.newArrayList();
for (ServerExpectation expectation : expectations) {
segmentIds.add(expectation.getSegmentId());
intervals.add(expectation.getInterval());
results.add(expectation.getResults());
}
EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context)))
.andReturn(toQueryableTimeseriesResults(expectBySegment, segmentIds, intervals, results))
.once();
} else if (query instanceof TopNQuery) {
List<String> segmentIds = Lists.newArrayList();
List<Interval> intervals = Lists.newArrayList();
List<Iterable<Result<TopNResultValue>>> results = Lists.newArrayList();
for (ServerExpectation expectation : expectations) {
segmentIds.add(expectation.getSegmentId());
intervals.add(expectation.getInterval());
results.add(expectation.getResults());
}
EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context)))
.andReturn(toQueryableTopNResults(segmentIds, intervals, results))
.once();
} else if (query instanceof SearchQuery) {
List<String> segmentIds = Lists.newArrayList();
List<Interval> intervals = Lists.newArrayList();
List<Iterable<Result<SearchResultValue>>> results = Lists.newArrayList();
for (ServerExpectation expectation : expectations) {
segmentIds.add(expectation.getSegmentId());
intervals.add(expectation.getInterval());
results.add(expectation.getResults());
}
EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context)))
.andReturn(toQueryableSearchResults(segmentIds, intervals, results))
.once();
} else if (query instanceof SelectQuery) {
List<String> segmentIds = Lists.newArrayList();
List<Interval> intervals = Lists.newArrayList();
List<Iterable<Result<SelectResultValue>>> results = Lists.newArrayList();
for (ServerExpectation expectation : expectations) {
segmentIds.add(expectation.getSegmentId());
intervals.add(expectation.getInterval());
results.add(expectation.getResults());
}
EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context)))
.andReturn(toQueryableSelectResults(segmentIds, intervals, results))
.once();
} else if (query instanceof GroupByQuery) {
List<String> segmentIds = Lists.newArrayList();
List<Interval> intervals = Lists.newArrayList();
List<Iterable<Row>> results = Lists.newArrayList();
for (ServerExpectation expectation : expectations) {
segmentIds.add(expectation.getSegmentId());
intervals.add(expectation.getInterval());
results.add(expectation.getResults());
}
EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context)))
.andReturn(toQueryableGroupByResults(segmentIds, intervals, results))
.once();
} else if (query instanceof TimeBoundaryQuery) {
List<String> segmentIds = Lists.newArrayList();
List<Interval> intervals = Lists.newArrayList();
List<Iterable<Result<TimeBoundaryResultValue>>> results = Lists.newArrayList();
for (ServerExpectation expectation : expectations) {
segmentIds.add(expectation.getSegmentId());
intervals.add(expectation.getInterval());
results.add(expectation.getResults());
}
EasyMock.expect(queryable.run(EasyMock.capture(capture), EasyMock.capture(context)))
.andReturn(toQueryableTimeBoundaryResults(segmentIds, intervals, results))
.once();
} else {
throw new ISE("Unknown query type[%s]", query.getClass());
}
}
final int expectedResultsRangeStart;
final int expectedResultsRangeEnd;