add more unit tests

processing/src/test/java/io/druid/query/RetryQueryRunnerTest.java

@@ -14,10 +14,12 @@ import io.druid.query.aggregation.LongSumAggregatorFactory;
 import io.druid.query.aggregation.MetricManipulationFn;
 import io.druid.query.timeseries.TimeseriesQuery;
 import io.druid.query.timeseries.TimeseriesResultValue;
+import io.druid.segment.SegmentMissingException;
 import org.joda.time.Interval;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -97,6 +99,94 @@ public class RetryQueryRunnerTest
           }
         },
         new RetryQueryRunnerConfig()
+        {
+          private int numTries = 0;
+          private boolean returnPartialResults = true;
+
+          public int numTries() { return numTries; }
+          public boolean returnPartialResults() { return returnPartialResults; }
+        }
+    );
+
+    Iterable<Result<TimeseriesResultValue>> actualResults = Sequences.toList(
+        runner.run(query, context),
+        Lists.<Result<TimeseriesResultValue>>newArrayList()
+    );
+
+    Assert.assertTrue("Should have one entry in the list of missing segments", ((List) context.get(RetryQueryRunner.missingSegments)).size() == 1);
+    Assert.assertTrue("Should return an empty sequence as a result", ((List) actualResults).size() == 0);
+  }
+
+  @Test
+  public void testRetry() throws Exception
+  {
+    Map<String, Object> context = new MapMaker().makeMap();
+    context.put("count", 0);
+    context.put("missingSegments", Lists.newArrayList());
+    RetryQueryRunner runner = new RetryQueryRunner(
+        new QueryRunner()
+        {
+          @Override
+          public Sequence run(Query query, Map context)
+          {
+            if (context.get("count") == 0) {
+              ((List) context.get(RetryQueryRunner.missingSegments)).add(
+                  new SegmentDescriptor(
+                      new Interval(
+                          178888,
+                          1999999
+                      ), "test", 1
+                  )
+              );
+              context.put("count", 1);
+              return Sequences.empty();
+            } else {
+              ArrayList lst = new ArrayList();
+              lst.add("hello world");
+              return Sequences.simple(lst);
+            }
+          }
+        },
+        new QueryToolChest()
+        {
+          @Override
+          public QueryRunner mergeResults(QueryRunner runner)
+          {
+            return null;
+          }
+
+          @Override
+          public Sequence mergeSequences(Sequence seqOfSequences)
+          {
+            return new OrderedMergeSequence<Result<TimeseriesResultValue>>(getOrdering(), seqOfSequences);
+          }
+
+          @Override
+          public ServiceMetricEvent.Builder makeMetricBuilder(Query query)
+          {
+            return null;
+          }
+
+          @Override
+          public Function makePreComputeManipulatorFn(
+              Query query, MetricManipulationFn fn
+          )
+          {
+            return null;
+          }
+
+          @Override
+          public TypeReference getResultTypeReference()
+          {
+            return null;
+          }
+
+          public Ordering<Result<TimeseriesResultValue>> getOrdering()
+          {
+            return Ordering.natural();
+          }
+        },
+        new RetryQueryRunnerConfig()
         {
           private int numTries = 1;
           private boolean returnPartialResults = true;
@@ -111,6 +201,85 @@ public class RetryQueryRunnerTest
         Lists.<Result<TimeseriesResultValue>>newArrayList()
     );
 
-    Assert.assertTrue("Should return an empty sequence as a result", ((List) actualResults).size() == 0);
+    actualResults = Sequences.toList(
+        runner.run(query, context),
+        Lists.<Result<TimeseriesResultValue>>newArrayList()
+    );
+
+    Assert.assertTrue("Should return a list with one element", ((List) actualResults).size() == 1);
+    Assert.assertTrue("Should have nothing in missingSegment list", ((List) context.get(RetryQueryRunner.missingSegments)).size() == 0);
   }
+
+  @Test(expected= SegmentMissingException.class)
+  public void testException() throws Exception
+  {
+    Map<String, Object> context = new MapMaker().makeMap();
+    context.put("missingSegments", Lists.newArrayList());
+    RetryQueryRunner runner = new RetryQueryRunner(
+        new QueryRunner()
+        {
+          @Override
+          public Sequence run(Query query, Map context)
+          {
+            ((List)context.get(RetryQueryRunner.missingSegments)).add(new SegmentDescriptor(new Interval(178888, 1999999), "test", 1));
+            return Sequences.empty();
+          }
+        },
+        new QueryToolChest()
+        {
+          @Override
+          public QueryRunner mergeResults(QueryRunner runner)
+          {
+            return null;
+          }
+
+          @Override
+          public Sequence mergeSequences(Sequence seqOfSequences)
+          {
+            return new OrderedMergeSequence<Result<TimeseriesResultValue>>(getOrdering(), seqOfSequences);
+          }
+
+          @Override
+          public ServiceMetricEvent.Builder makeMetricBuilder(Query query)
+          {
+            return null;
+          }
+
+          @Override
+          public Function makePreComputeManipulatorFn(
+              Query query, MetricManipulationFn fn
+          )
+          {
+            return null;
+          }
+
+          @Override
+          public TypeReference getResultTypeReference()
+          {
+            return null;
+          }
+
+          public Ordering<Result<TimeseriesResultValue>> getOrdering()
+          {
+            return Ordering.natural();
+          }
+        },
+        new RetryQueryRunnerConfig()
+        {
+          private int numTries = 1;
+          private boolean returnPartialResults = false;
+
+          public int numTries() { return numTries; }
+          public boolean returnPartialResults() { return returnPartialResults; }
+        }
+    );
+
+    Iterable<Result<TimeseriesResultValue>> actualResults = Sequences.toList(
+        runner.run(query, context),
+        Lists.<Result<TimeseriesResultValue>>newArrayList()
+    );
+
+    Assert.assertTrue("Should have one entry in the list of missing segments", ((List) context.get(RetryQueryRunner.missingSegments)).size() == 1);
+  }
+
 }

processing/src/test/java/io/druid/query/timeseries/TimeseriesQueryRunnerTest.java

@@ -252,20 +252,20 @@ public class TimeseriesQueryRunnerTest
   public void testTimeseries()
   {
     TimeseriesQuery query = Druids.newTimeseriesQueryBuilder()
-                                  .dataSource(QueryRunnerTestHelper.dataSource)
+                                                                         .dataSource(QueryRunnerTestHelper.dataSource)
-                                  .granularity(QueryRunnerTestHelper.dayGran)
+                                                                         .granularity(QueryRunnerTestHelper.dayGran)
-                                  .intervals(QueryRunnerTestHelper.firstToThird)
+                                                                         .intervals(QueryRunnerTestHelper.firstToThird)
-                                  .aggregators(
+                                                                         .aggregators(
-                                      Arrays.<AggregatorFactory>asList(
+                                                                             Arrays.<AggregatorFactory>asList(
-                                          QueryRunnerTestHelper.rowsCount,
+                                                                                 QueryRunnerTestHelper.rowsCount,
-                                          new LongSumAggregatorFactory(
+                                                                                 new LongSumAggregatorFactory(
-                                              "idx",
+                                                                                     "idx",
-                                              "index"
+                                                                                     "index"
-                                          ),
+                                                                                 ),
-                                          QueryRunnerTestHelper.qualityUniques
+                                                                                 QueryRunnerTestHelper.qualityUniques
-                                      )
+                                                                             )
-                                  )
+                                                                         )
-                                  .build();
+                                                                         .build();
 
     List<Result<TimeseriesResultValue>> expectedResults = Arrays.asList(
         new Result<TimeseriesResultValue>(

server/src/main/java/io/druid/client/CachingClusteredClient.java

@@ -253,7 +253,6 @@ public class CachingClusteredClient<T> implements QueryRunner<T>
             addSequencesFromServer(listOfSequences);
             addSequencesFromCache(listOfSequences);
 
-
             Collections.sort(
                 listOfSequences,
                 Ordering.natural().onResultOf(Pair.<DateTime, Sequence<T>>lhsFn())