From a76bf9ab2ac773c6a80814f9deec57bbe287558c Mon Sep 17 00:00:00 2001
From: Himanshu <g.himanshu@gmail.com>
Date: Wed, 22 Aug 2018 16:38:36 -0700
Subject: [PATCH] add ability to do optional rollup in AggregationTestHelper
 (#6213)

---
 .../aggregation/AggregationTestHelper.java    | 82 +++++++++++++++++--
 1 file changed, 75 insertions(+), 7 deletions(-)

diff --git a/processing/src/test/java/io/druid/query/aggregation/AggregationTestHelper.java b/processing/src/test/java/io/druid/query/aggregation/AggregationTestHelper.java
index 869d1acdc3a..9ab3b5fb47b 100644
--- a/processing/src/test/java/io/druid/query/aggregation/AggregationTestHelper.java
+++ b/processing/src/test/java/io/druid/query/aggregation/AggregationTestHelper.java
@@ -343,7 +343,23 @@ public class AggregationTestHelper implements Closeable
   ) throws Exception
   {
     File segmentDir = tempFolder.newFolder();
-    createIndex(inputDataFile, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
+    createIndex(inputDataFile, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount, true);
+    return runQueryOnSegments(Lists.newArrayList(segmentDir), groupByQueryJson);
+  }
+
+  public Sequence<Row> createIndexAndRunQueryOnSegment(
+      File inputDataFile,
+      String parserJson,
+      String aggregators,
+      long minTimestamp,
+      Granularity gran,
+      int maxRowCount,
+      boolean rollup,
+      String groupByQueryJson
+  ) throws Exception
+  {
+    File segmentDir = tempFolder.newFolder();
+    createIndex(inputDataFile, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount, rollup);
     return runQueryOnSegments(Collections.singletonList(segmentDir), groupByQueryJson);
   }
 
@@ -356,9 +372,32 @@ public class AggregationTestHelper implements Closeable
       int maxRowCount,
       String groupByQueryJson
   ) throws Exception
+  {
+    return createIndexAndRunQueryOnSegment(
+        inputDataStream,
+        parserJson,
+        aggregators,
+        minTimestamp,
+        gran,
+        maxRowCount,
+        true,
+        groupByQueryJson
+    );
+  }
+
+  public Sequence<Row> createIndexAndRunQueryOnSegment(
+      InputStream inputDataStream,
+      String parserJson,
+      String aggregators,
+      long minTimestamp,
+      Granularity gran,
+      int maxRowCount,
+      boolean rollup,
+      String groupByQueryJson
+  ) throws Exception
   {
     File segmentDir = tempFolder.newFolder();
-    createIndex(inputDataStream, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount);
+    createIndex(inputDataStream, parserJson, aggregators, segmentDir, minTimestamp, gran, maxRowCount, rollup);
     return runQueryOnSegments(Collections.singletonList(segmentDir), groupByQueryJson);
   }
 
@@ -379,7 +418,31 @@ public class AggregationTestHelper implements Closeable
         outDir,
         minTimestamp,
         gran,
-        maxRowCount
+        maxRowCount,
+        true
+    );
+  }
+
+  public void createIndex(
+      File inputDataFile,
+      String parserJson,
+      String aggregators,
+      File outDir,
+      long minTimestamp,
+      Granularity gran,
+      int maxRowCount,
+      boolean rollup
+  ) throws Exception
+  {
+    createIndex(
+        new FileInputStream(inputDataFile),
+        parserJson,
+        aggregators,
+        outDir,
+        minTimestamp,
+        gran,
+        maxRowCount,
+        rollup
     );
   }
 
@@ -390,7 +453,8 @@ public class AggregationTestHelper implements Closeable
       File outDir,
       long minTimestamp,
       Granularity gran,
-      int maxRowCount
+      int maxRowCount,
+      boolean rollup
   ) throws Exception
   {
     try {
@@ -412,7 +476,8 @@ public class AggregationTestHelper implements Closeable
           minTimestamp,
           gran,
           true,
-          maxRowCount
+          maxRowCount,
+          rollup
       );
     }
     finally {
@@ -428,7 +493,8 @@ public class AggregationTestHelper implements Closeable
       long minTimestamp,
       Granularity gran,
       boolean deserializeComplexMetrics,
-      int maxRowCount
+      int maxRowCount,
+      boolean rollup
   ) throws Exception
   {
     IncrementalIndex index = null;
@@ -441,6 +507,7 @@ public class AggregationTestHelper implements Closeable
                   .withMinTimestamp(minTimestamp)
                   .withQueryGranularity(gran)
                   .withMetrics(metrics)
+                  .withRollup(rollup)
                   .build()
           )
           .setDeserializeComplexMetrics(deserializeComplexMetrics)
@@ -460,6 +527,7 @@ public class AggregationTestHelper implements Closeable
                       .withMinTimestamp(minTimestamp)
                       .withQueryGranularity(gran)
                       .withMetrics(metrics)
+                      .withRollup(rollup)
                       .build()
               )
               .setDeserializeComplexMetrics(deserializeComplexMetrics)
@@ -484,7 +552,7 @@ public class AggregationTestHelper implements Closeable
         for (File file : toMerge) {
           indexes.add(indexIO.loadIndex(file));
         }
-        indexMerger.mergeQueryableIndex(indexes, true, metrics, outDir, new IndexSpec(), null);
+        indexMerger.mergeQueryableIndex(indexes, rollup, metrics, outDir, new IndexSpec(), null);
 
         for (QueryableIndex qi : indexes) {
           qi.close();