Improve agg reduce tests

test/framework/src/main/java/org/elasticsearch/test/InternalAggregationTestCase.java

@@ -281,24 +281,34 @@ protected T createUnmappedInstance(String name, Map<String, Object> metadata) {
         return createTestInstance(name, metadata);
     }
 
-    public void testReduceRandom() throws IOException {
-        String name = randomAlphaOfLength(5);
+    /**
+     * Generate a list of inputs to reduce. Defaults to calling
+     * {@link #createTestInstance(String)} and
+     * {@link #createUnmappedInstance(String)} but should be overridden
+     * if it isn't realistic to reduce test instances.
+     */
+    protected List<T> randomResultsToReduce(String name, int size) {
         List<T> inputs = new ArrayList<>();
-        List<InternalAggregation> toReduce = new ArrayList<>();
-        int toReduceSize = between(1, 200);
-        for (int i = 0; i < toReduceSize; i++) {
+        for (int i = 0; i < size; i++) {
             T t = randomBoolean() ? createUnmappedInstance(name) : createTestInstance(name);
             inputs.add(t);
-            toReduce.add(t);
         }
+        return inputs;
+    }
+
+    public void testReduceRandom() throws IOException {
+        String name = randomAlphaOfLength(5);
+        List<T> inputs = randomResultsToReduce(name, between(1, 200));
+        List<InternalAggregation> toReduce = new ArrayList<>();
+        toReduce.addAll(inputs);
         // Sort aggs so that unmapped come last.  This mimicks the behavior of InternalAggregations.reduce()
         inputs.sort(INTERNAL_AGG_COMPARATOR);
         ScriptService mockScriptService = mockScriptService();
         MockBigArrays bigArrays = new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), new NoneCircuitBreakerService());
         if (randomBoolean() && toReduce.size() > 1) {
             // sometimes do a partial reduce
             Collections.shuffle(toReduce, random());
-            int r = randomIntBetween(1, toReduceSize);
+            int r = randomIntBetween(1, inputs.size());
             List<InternalAggregation> internalAggregations = toReduce.subList(0, r);
             InternalAggregation.ReduceContext context = InternalAggregation.ReduceContext.forPartialReduction(
                     bigArrays, mockScriptService, () -> PipelineAggregator.PipelineTree.EMPTY);
@@ -319,7 +329,7 @@ public void testReduceRandom() throws IOException {
             if (randomBoolean()) {
                 reduced = copyInstance(reduced);
             }
-            toReduce = new ArrayList<>(toReduce.subList(r, toReduceSize));
+            toReduce = new ArrayList<>(toReduce.subList(r, inputs.size()));
             toReduce.add(reduced);
         }
         MultiBucketConsumer bucketConsumer = new MultiBucketConsumer(DEFAULT_MAX_BUCKETS,

x-pack/plugin/analytics/src/test/java/org/elasticsearch/xpack/analytics/ttest/InternalTTestTests.java

@@ -37,27 +37,32 @@ public class InternalTTestTests extends InternalAggregationTestCase<InternalTTes
 
     @Override
     protected InternalTTest createTestInstance(String name, Map<String, Object> metadata) {
-        TTestState state = randomState();
+        TTestState state = randomState(Long.MAX_VALUE);
         DocValueFormat formatter = randomNumericDocValueFormat();
         return new InternalTTest(name, state, formatter, metadata);
     }
 
-    private TTestState randomState() {
+    @Override
+    protected List<InternalTTest> randomResultsToReduce(String name, int size) {
+        List<InternalTTest> inputs = new ArrayList<>(size);
+        for (int i = 0; i < size; i++) {
+            TTestState state = randomState(Long.MAX_VALUE / size); // Make sure the sum of all the counts doesn't wrap
+            DocValueFormat formatter = randomNumericDocValueFormat();
+            inputs.add(new InternalTTest(name, state, formatter, null));
+        }
+        return inputs;
+    }
+
+    private TTestState randomState(long maxCount) {
         if (type == TTestType.PAIRED) {
-            return new PairedTTestState(randomStats(), tails);
+            return new PairedTTestState(randomStats(maxCount), tails);
         } else {
-            return new UnpairedTTestState(randomStats(), randomStats(), type == TTestType.HOMOSCEDASTIC, tails);
+            return new UnpairedTTestState(randomStats(maxCount), randomStats(maxCount), type == TTestType.HOMOSCEDASTIC, tails);
         }
     }
 
-    private TTestStats randomStats() {
-        /*
-         * Use a count significantly less than Long.MAX_VALUE so the reduce
-         * phase doesn't wrap to a negative number. If it *did* then we'd
-         * try to serialize a negative number with writeVLong which throws
-         * an assertion in tests.
-         */
-        return new TTestStats(randomLongBetween(0, Integer.MAX_VALUE), randomDouble(), randomDouble());
+    private TTestStats randomStats(long maxCount) {
+        return new TTestStats(randomLongBetween(0, maxCount), randomDouble(), randomDouble());
     }
 
     @Override
@@ -97,7 +102,7 @@ protected InternalTTest mutateInstance(InternalTTest instance) {
                 name += randomAlphaOfLength(5);
                 break;
             case 1:
-                state = randomState();
+                state = randomState(Long.MAX_VALUE);
                 break;
             case 2:
                 if (metadata == null) {