Remove explicit SearchResponse references from server bucket aggs (part4) (#102038)

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/RandomSamplerIT.java

@@ -9,8 +9,6 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.elasticsearch.action.index.IndexRequestBuilder;
-import org.elasticsearch.action.search.SearchRequest;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;
 import org.elasticsearch.search.aggregations.bucket.sampler.random.InternalRandomSampler;
 import org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplerAggregationBuilder;
@@ -24,6 +22,7 @@ import java.util.Map;
 
 import static org.elasticsearch.search.aggregations.AggregationBuilders.avg;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.histogram;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.hamcrest.Matchers.lessThan;
 
@@ -87,40 +86,47 @@ public class RandomSamplerIT extends ESIntegTestCase {
     }
 
     public void testRandomSampler() {
-        double sampleMonotonicValue = 0.0;
-        double sampleNumericValue = 0.0;
-        double sampledDocCount = 0.0;
+        double[] sampleMonotonicValue = new double[1];
+        double[] sampleNumericValue = new double[1];
+        double[] sampledDocCount = new double[1];
 
         for (int i = 0; i < NUM_SAMPLE_RUNS; i++) {
-            SearchRequest sampledRequest = prepareSearch("idx").addAggregation(
-                new RandomSamplerAggregationBuilder("sampler").setProbability(PROBABILITY)
-                    .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
-                    .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
-            ).request();
-            InternalRandomSampler sampler = client().search(sampledRequest).actionGet().getAggregations().get("sampler");
-            sampleMonotonicValue += ((Avg) sampler.getAggregations().get("mean_monotonic")).getValue();
-            sampleNumericValue += ((Avg) sampler.getAggregations().get("mean_numeric")).getValue();
-            sampledDocCount += sampler.getDocCount();
+            assertResponse(
+                prepareSearch("idx").addAggregation(
+                    new RandomSamplerAggregationBuilder("sampler").setProbability(PROBABILITY)
+                        .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
+                        .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
+                ),
+                response -> {
+                    InternalRandomSampler sampler = response.getAggregations().get("sampler");
+                    sampleMonotonicValue[0] += ((Avg) sampler.getAggregations().get("mean_monotonic")).getValue();
+                    sampleNumericValue[0] += ((Avg) sampler.getAggregations().get("mean_numeric")).getValue();
+                    sampledDocCount[0] += sampler.getDocCount();
+                }
+            );
         }
-        sampledDocCount /= NUM_SAMPLE_RUNS;
-        sampleMonotonicValue /= NUM_SAMPLE_RUNS;
-        sampleNumericValue /= NUM_SAMPLE_RUNS;
+        sampledDocCount[0] /= NUM_SAMPLE_RUNS;
+        sampleMonotonicValue[0] /= NUM_SAMPLE_RUNS;
+        sampleNumericValue[0] /= NUM_SAMPLE_RUNS;
         double expectedDocCount = PROBABILITY * numDocs;
         // We're taking the mean of NUM_SAMPLE_RUNS for which each run has standard deviation
         // sqrt(PROBABILITY * numDocs) so the 6 sigma error, for which we expect 1 failure in
         // 500M runs, is 6 * sqrt(PROBABILITY * numDocs / NUM_SAMPLE_RUNS).
         double maxCountError = 6.0 * Math.sqrt(PROBABILITY * numDocs / NUM_SAMPLE_RUNS);
-        assertThat(Math.abs(sampledDocCount - expectedDocCount), lessThan(maxCountError));
-
-        SearchResponse trueValueResponse = prepareSearch("idx").addAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
-            .addAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
-            .get();
-        double trueMonotonic = ((Avg) trueValueResponse.getAggregations().get("mean_monotonic")).getValue();
-        double trueNumeric = ((Avg) trueValueResponse.getAggregations().get("mean_numeric")).getValue();
-        double maxMonotonicError = 6.0 * Math.sqrt(varMonotonic / (numDocs * PROBABILITY * NUM_SAMPLE_RUNS));
-        double maxNumericError = 6.0 * Math.sqrt(varNumeric / (numDocs * PROBABILITY * NUM_SAMPLE_RUNS));
-        assertThat(Math.abs(sampleMonotonicValue - trueMonotonic), lessThan(maxMonotonicError));
-        assertThat(Math.abs(sampleNumericValue - trueNumeric), lessThan(maxNumericError));
+        assertThat(Math.abs(sampledDocCount[0] - expectedDocCount), lessThan(maxCountError));
+
+        assertResponse(
+            prepareSearch("idx").addAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
+                .addAggregation(avg("mean_numeric").field(NUMERIC_VALUE)),
+            response -> {
+                double trueMonotonic = ((Avg) response.getAggregations().get("mean_monotonic")).getValue();
+                double trueNumeric = ((Avg) response.getAggregations().get("mean_numeric")).getValue();
+                double maxMonotonicError = 6.0 * Math.sqrt(varMonotonic / (numDocs * PROBABILITY * NUM_SAMPLE_RUNS));
+                double maxNumericError = 6.0 * Math.sqrt(varNumeric / (numDocs * PROBABILITY * NUM_SAMPLE_RUNS));
+                assertThat(Math.abs(sampleMonotonicValue[0] - trueMonotonic), lessThan(maxMonotonicError));
+                assertThat(Math.abs(sampleNumericValue[0] - trueNumeric), lessThan(maxNumericError));
+            }
+        );
     }
 
     public void testRandomSamplerHistogram() {
@@ -129,28 +135,32 @@ public class RandomSamplerIT extends ESIntegTestCase {
         Map<String, Double> sampledDocCount = new HashMap<>();
 
         for (int i = 0; i < NUM_SAMPLE_RUNS; i++) {
-            SearchRequest sampledRequest = prepareSearch("idx").addAggregation(
-                new RandomSamplerAggregationBuilder("sampler").setProbability(PROBABILITY)
-                    .subAggregation(
-                        histogram("histo").field(NUMERIC_VALUE)
-                            .interval(5.0)
-                            .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
-                            .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
-                    )
-            ).request();
-            InternalRandomSampler sampler = client().search(sampledRequest).actionGet().getAggregations().get("sampler");
-            Histogram histo = sampler.getAggregations().get("histo");
-            for (Histogram.Bucket bucket : histo.getBuckets()) {
-                sampleMonotonicValue.compute(
-                    bucket.getKeyAsString(),
-                    (k, v) -> ((Avg) bucket.getAggregations().get("mean_monotonic")).getValue() + (v == null ? 0 : v)
-                );
-                sampleNumericValue.compute(
-                    bucket.getKeyAsString(),
-                    (k, v) -> ((Avg) bucket.getAggregations().get("mean_numeric")).getValue() + (v == null ? 0 : v)
-                );
-                sampledDocCount.compute(bucket.getKeyAsString(), (k, v) -> bucket.getDocCount() + (v == null ? 0 : v));
-            }
+            assertResponse(
+                prepareSearch("idx").addAggregation(
+                    new RandomSamplerAggregationBuilder("sampler").setProbability(PROBABILITY)
+                        .subAggregation(
+                            histogram("histo").field(NUMERIC_VALUE)
+                                .interval(5.0)
+                                .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
+                                .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
+                        )
+                ),
+                response -> {
+                    InternalRandomSampler sampler = response.getAggregations().get("sampler");
+                    Histogram histo = sampler.getAggregations().get("histo");
+                    for (Histogram.Bucket bucket : histo.getBuckets()) {
+                        sampleMonotonicValue.compute(
+                            bucket.getKeyAsString(),
+                            (k, v) -> ((Avg) bucket.getAggregations().get("mean_monotonic")).getValue() + (v == null ? 0 : v)
+                        );
+                        sampleNumericValue.compute(
+                            bucket.getKeyAsString(),
+                            (k, v) -> ((Avg) bucket.getAggregations().get("mean_numeric")).getValue() + (v == null ? 0 : v)
+                        );
+                        sampledDocCount.compute(bucket.getKeyAsString(), (k, v) -> bucket.getDocCount() + (v == null ? 0 : v));
+                    }
+                }
+            );
         }
         for (String key : sampledDocCount.keySet()) {
             sampledDocCount.put(key, sampledDocCount.get(key) / NUM_SAMPLE_RUNS);
@@ -158,25 +168,29 @@ public class RandomSamplerIT extends ESIntegTestCase {
             sampleMonotonicValue.put(key, sampleMonotonicValue.get(key) / NUM_SAMPLE_RUNS);
         }
 
-        SearchResponse trueValueResponse = prepareSearch("idx").addAggregation(
-            histogram("histo").field(NUMERIC_VALUE)
-                .interval(5.0)
-                .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
-                .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
-        ).get();
-        Histogram histogram = trueValueResponse.getAggregations().get("histo");
-        for (Histogram.Bucket bucket : histogram.getBuckets()) {
-            long numDocs = bucket.getDocCount();
-            // Note the true count is estimated by dividing the bucket sample doc count by PROBABILITY.
-            double maxCountError = 6.0 * Math.sqrt(numDocs / NUM_SAMPLE_RUNS / (0.5 * PROBABILITY));
-            assertThat(Math.abs(sampledDocCount.get(bucket.getKeyAsString()) - numDocs), lessThan(maxCountError));
-            double trueMonotonic = ((Avg) bucket.getAggregations().get("mean_monotonic")).getValue();
-            double trueNumeric = ((Avg) bucket.getAggregations().get("mean_numeric")).getValue();
-            double maxMonotonicError = 6.0 * Math.sqrt(varMonotonic / (numDocs * 0.5 * PROBABILITY * NUM_SAMPLE_RUNS));
-            double maxNumericError = 6.0 * Math.sqrt(varNumeric / (numDocs * 0.5 * PROBABILITY * NUM_SAMPLE_RUNS));
-            assertThat(Math.abs(sampleMonotonicValue.get(bucket.getKeyAsString()) - trueMonotonic), lessThan(maxMonotonicError));
-            assertThat(Math.abs(sampleNumericValue.get(bucket.getKeyAsString()) - trueNumeric), lessThan(maxNumericError));
-        }
+        assertResponse(
+            prepareSearch("idx").addAggregation(
+                histogram("histo").field(NUMERIC_VALUE)
+                    .interval(5.0)
+                    .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))
+                    .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))
+            ),
+            response -> {
+                Histogram histogram = response.getAggregations().get("histo");
+                for (Histogram.Bucket bucket : histogram.getBuckets()) {
+                    long numDocs = bucket.getDocCount();
+                    // Note the true count is estimated by dividing the bucket sample doc count by PROBABILITY.
+                    double maxCountError = 6.0 * Math.sqrt(numDocs / NUM_SAMPLE_RUNS / (0.5 * PROBABILITY));
+                    assertThat(Math.abs(sampledDocCount.get(bucket.getKeyAsString()) - numDocs), lessThan(maxCountError));
+                    double trueMonotonic = ((Avg) bucket.getAggregations().get("mean_monotonic")).getValue();
+                    double trueNumeric = ((Avg) bucket.getAggregations().get("mean_numeric")).getValue();
+                    double maxMonotonicError = 6.0 * Math.sqrt(varMonotonic / (numDocs * 0.5 * PROBABILITY * NUM_SAMPLE_RUNS));
+                    double maxNumericError = 6.0 * Math.sqrt(varNumeric / (numDocs * 0.5 * PROBABILITY * NUM_SAMPLE_RUNS));
+                    assertThat(Math.abs(sampleMonotonicValue.get(bucket.getKeyAsString()) - trueMonotonic), lessThan(maxMonotonicError));
+                    assertThat(Math.abs(sampleNumericValue.get(bucket.getKeyAsString()) - trueNumeric), lessThan(maxNumericError));
+                }
+            }
+        );
     }
 
 }

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/ReverseNestedIT.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.elasticsearch.action.search.SearchPhaseExecutionException;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.search.aggregations.Aggregator.SubAggCollectionMode;
 import org.elasticsearch.search.aggregations.BucketOrder;
 import org.elasticsearch.search.aggregations.InternalAggregation;
@@ -34,7 +33,7 @@ import static org.elasticsearch.search.aggregations.AggregationBuilders.reverseN
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertHitCount;
-import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
@@ -143,310 +142,314 @@ public class ReverseNestedIT extends ESIntegTestCase {
     }
 
     public void testSimpleReverseNestedToRoot() throws Exception {
-        SearchResponse response = prepareSearch("idx1").addAggregation(
-            nested("nested1", "nested1").subAggregation(
-                terms("field2").field("nested1.field2")
-                    .subAggregation(
-                        reverseNested("nested1_to_field1").subAggregation(
-                            terms("field1").field("field1").collectMode(randomFrom(SubAggCollectionMode.values()))
+        assertNoFailuresAndResponse(
+            prepareSearch("idx1").addAggregation(
+                nested("nested1", "nested1").subAggregation(
+                    terms("field2").field("nested1.field2")
+                        .subAggregation(
+                            reverseNested("nested1_to_field1").subAggregation(
+                                terms("field1").field("field1").collectMode(randomFrom(SubAggCollectionMode.values()))
+                            )
                         )
-                    )
-            )
-        ).get();
-
-        assertNoFailures(response);
-
-        Nested nested = response.getAggregations().get("nested1");
-        assertThat(nested, notNullValue());
-        assertThat(nested.getName(), equalTo("nested1"));
-        assertThat(nested.getDocCount(), equalTo(25L));
-        assertThat(nested.getAggregations().asList().isEmpty(), is(false));
-
-        Terms usernames = nested.getAggregations().get("field2");
-        assertThat(usernames, notNullValue());
-        assertThat(usernames.getBuckets().size(), equalTo(9));
-        List<Terms.Bucket> usernameBuckets = new ArrayList<>(usernames.getBuckets());
-
-        // nested.field2: 1
-        Terms.Bucket bucket = usernameBuckets.get(0);
-        assertThat(bucket.getKeyAsString(), equalTo("1"));
-        assertThat(bucket.getDocCount(), equalTo(6L));
-        ReverseNested reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        assertThat(((InternalAggregation) reverseNested).getProperty("_count"), equalTo(5L));
-        Terms tags = reverseNested.getAggregations().get("field1");
-        assertThat(((InternalAggregation) reverseNested).getProperty("field1"), sameInstance(tags));
-        List<Terms.Bucket> tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(6));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(4L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(3L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(4).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(4).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(5).getKeyAsString(), equalTo("x"));
-        assertThat(tagsBuckets.get(5).getDocCount(), equalTo(1L));
-
-        // nested.field2: 4
-        bucket = usernameBuckets.get(1);
-        assertThat(bucket.getKeyAsString(), equalTo("4"));
-        assertThat(bucket.getDocCount(), equalTo(4L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(5));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(3L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(4).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(4).getDocCount(), equalTo(1L));
-
-        // nested.field2: 7
-        bucket = usernameBuckets.get(2);
-        assertThat(bucket.getKeyAsString(), equalTo("7"));
-        assertThat(bucket.getDocCount(), equalTo(3L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(5));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(4).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(4).getDocCount(), equalTo(1L));
-
-        // nested.field2: 2
-        bucket = usernameBuckets.get(3);
-        assertThat(bucket.getKeyAsString(), equalTo("2"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(3));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-
-        // nested.field2: 3
-        bucket = usernameBuckets.get(4);
-        assertThat(bucket.getKeyAsString(), equalTo("3"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(3));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-
-        // nested.field2: 5
-        bucket = usernameBuckets.get(5);
-        assertThat(bucket.getKeyAsString(), equalTo("5"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(4));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("z"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-
-        // nested.field2: 6
-        bucket = usernameBuckets.get(6);
-        assertThat(bucket.getKeyAsString(), equalTo("6"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(4));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("y"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-
-        // nested.field2: 8
-        bucket = usernameBuckets.get(7);
-        assertThat(bucket.getKeyAsString(), equalTo("8"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(4));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("x"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-
-        // nested.field2: 9
-        bucket = usernameBuckets.get(8);
-        assertThat(bucket.getKeyAsString(), equalTo("9"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(4));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("z"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+                )
+            ),
+            response -> {
+                Nested nested = response.getAggregations().get("nested1");
+                assertThat(nested, notNullValue());
+                assertThat(nested.getName(), equalTo("nested1"));
+                assertThat(nested.getDocCount(), equalTo(25L));
+                assertThat(nested.getAggregations().asList().isEmpty(), is(false));
+
+                Terms usernames = nested.getAggregations().get("field2");
+                assertThat(usernames, notNullValue());
+                assertThat(usernames.getBuckets().size(), equalTo(9));
+                List<Terms.Bucket> usernameBuckets = new ArrayList<>(usernames.getBuckets());
+
+                // nested.field2: 1
+                Terms.Bucket bucket = usernameBuckets.get(0);
+                assertThat(bucket.getKeyAsString(), equalTo("1"));
+                assertThat(bucket.getDocCount(), equalTo(6L));
+                ReverseNested reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                assertThat(((InternalAggregation) reverseNested).getProperty("_count"), equalTo(5L));
+                Terms tags = reverseNested.getAggregations().get("field1");
+                assertThat(((InternalAggregation) reverseNested).getProperty("field1"), sameInstance(tags));
+                List<Terms.Bucket> tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(6));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(4L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(3L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(4).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(4).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(5).getKeyAsString(), equalTo("x"));
+                assertThat(tagsBuckets.get(5).getDocCount(), equalTo(1L));
+
+                // nested.field2: 4
+                bucket = usernameBuckets.get(1);
+                assertThat(bucket.getKeyAsString(), equalTo("4"));
+                assertThat(bucket.getDocCount(), equalTo(4L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(5));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(3L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(4).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(4).getDocCount(), equalTo(1L));
+
+                // nested.field2: 7
+                bucket = usernameBuckets.get(2);
+                assertThat(bucket.getKeyAsString(), equalTo("7"));
+                assertThat(bucket.getDocCount(), equalTo(3L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(5));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(4).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(4).getDocCount(), equalTo(1L));
+
+                // nested.field2: 2
+                bucket = usernameBuckets.get(3);
+                assertThat(bucket.getKeyAsString(), equalTo("2"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(3));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+
+                // nested.field2: 3
+                bucket = usernameBuckets.get(4);
+                assertThat(bucket.getKeyAsString(), equalTo("3"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(3));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+
+                // nested.field2: 5
+                bucket = usernameBuckets.get(5);
+                assertThat(bucket.getKeyAsString(), equalTo("5"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(4));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("z"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+
+                // nested.field2: 6
+                bucket = usernameBuckets.get(6);
+                assertThat(bucket.getKeyAsString(), equalTo("6"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(4));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("y"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+
+                // nested.field2: 8
+                bucket = usernameBuckets.get(7);
+                assertThat(bucket.getKeyAsString(), equalTo("8"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(4));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(2L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("x"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+
+                // nested.field2: 9
+                bucket = usernameBuckets.get(8);
+                assertThat(bucket.getKeyAsString(), equalTo("9"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(4));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("z"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+            }
+        );
     }
 
     public void testSimpleNested1ToRootToNested2() throws Exception {
-        SearchResponse response = prepareSearch("idx2").addAggregation(
-            nested("nested1", "nested1").subAggregation(
-                reverseNested("nested1_to_root").subAggregation(nested("root_to_nested2", "nested1.nested2"))
-            )
-        ).get();
-
-        assertNoFailures(response);
-        Nested nested = response.getAggregations().get("nested1");
-        assertThat(nested.getName(), equalTo("nested1"));
-        assertThat(nested.getDocCount(), equalTo(9L));
-        ReverseNested reverseNested = nested.getAggregations().get("nested1_to_root");
-        assertThat(reverseNested.getName(), equalTo("nested1_to_root"));
-        assertThat(reverseNested.getDocCount(), equalTo(4L));
-        nested = reverseNested.getAggregations().get("root_to_nested2");
-        assertThat(nested.getName(), equalTo("root_to_nested2"));
-        assertThat(nested.getDocCount(), equalTo(27L));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx2").addAggregation(
+                nested("nested1", "nested1").subAggregation(
+                    reverseNested("nested1_to_root").subAggregation(nested("root_to_nested2", "nested1.nested2"))
+                )
+            ),
+            response -> {
+                Nested nested = response.getAggregations().get("nested1");
+                assertThat(nested.getName(), equalTo("nested1"));
+                assertThat(nested.getDocCount(), equalTo(9L));
+                ReverseNested reverseNested = nested.getAggregations().get("nested1_to_root");
+                assertThat(reverseNested.getName(), equalTo("nested1_to_root"));
+                assertThat(reverseNested.getDocCount(), equalTo(4L));
+                nested = reverseNested.getAggregations().get("root_to_nested2");
+                assertThat(nested.getName(), equalTo("root_to_nested2"));
+                assertThat(nested.getDocCount(), equalTo(27L));
+            }
+        );
     }
 
     public void testSimpleReverseNestedToNested1() throws Exception {
-        SearchResponse response = prepareSearch("idx2").addAggregation(
-            nested("nested1", "nested1.nested2").subAggregation(
-                terms("field2").field("nested1.nested2.field2")
-                    .order(BucketOrder.key(true))
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .size(10000)
-                    .subAggregation(
-                        reverseNested("nested1_to_field1").path("nested1")
-                            .subAggregation(
-                                terms("field1").field("nested1.field1")
-                                    .order(BucketOrder.key(true))
-                                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                            )
-                    )
-            )
-        ).get();
-
-        assertNoFailures(response);
-
-        Nested nested = response.getAggregations().get("nested1");
-        assertThat(nested, notNullValue());
-        assertThat(nested.getName(), equalTo("nested1"));
-        assertThat(nested.getDocCount(), equalTo(27L));
-        assertThat(nested.getAggregations().asList().isEmpty(), is(false));
-
-        Terms usernames = nested.getAggregations().get("field2");
-        assertThat(usernames, notNullValue());
-        assertThat(usernames.getBuckets().size(), equalTo(5));
-        List<Terms.Bucket> usernameBuckets = new ArrayList<>(usernames.getBuckets());
-
-        Terms.Bucket bucket = usernameBuckets.get(0);
-        assertThat(bucket.getKeyAsString(), equalTo("0"));
-        assertThat(bucket.getDocCount(), equalTo(12L));
-        ReverseNested reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        assertThat(reverseNested.getDocCount(), equalTo(5L));
-        Terms tags = reverseNested.getAggregations().get("field1");
-        List<Terms.Bucket> tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(2));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(3L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
-
-        bucket = usernameBuckets.get(1);
-        assertThat(bucket.getKeyAsString(), equalTo("1"));
-        assertThat(bucket.getDocCount(), equalTo(6L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        assertThat(reverseNested.getDocCount(), equalTo(4L));
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(4));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-
-        bucket = usernameBuckets.get(2);
-        assertThat(bucket.getKeyAsString(), equalTo("2"));
-        assertThat(bucket.getDocCount(), equalTo(5L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        assertThat(reverseNested.getDocCount(), equalTo(4L));
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(4));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
-        assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
-        assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("e"));
-        assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
-
-        bucket = usernameBuckets.get(3);
-        assertThat(bucket.getKeyAsString(), equalTo("3"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        assertThat(reverseNested.getDocCount(), equalTo(2L));
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(2));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("f"));
-
-        bucket = usernameBuckets.get(4);
-        assertThat(bucket.getKeyAsString(), equalTo("4"));
-        assertThat(bucket.getDocCount(), equalTo(2L));
-        reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        assertThat(reverseNested.getDocCount(), equalTo(2L));
-        tags = reverseNested.getAggregations().get("field1");
-        tagsBuckets = new ArrayList<>(tags.getBuckets());
-        assertThat(tagsBuckets.size(), equalTo(2));
-        assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("d"));
-        assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
-        assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("f"));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx2").addAggregation(
+                nested("nested1", "nested1.nested2").subAggregation(
+                    terms("field2").field("nested1.nested2.field2")
+                        .order(BucketOrder.key(true))
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .size(10000)
+                        .subAggregation(
+                            reverseNested("nested1_to_field1").path("nested1")
+                                .subAggregation(
+                                    terms("field1").field("nested1.field1")
+                                        .order(BucketOrder.key(true))
+                                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                                )
+                        )
+                )
+            ),
+            response -> {
+                Nested nested = response.getAggregations().get("nested1");
+                assertThat(nested, notNullValue());
+                assertThat(nested.getName(), equalTo("nested1"));
+                assertThat(nested.getDocCount(), equalTo(27L));
+                assertThat(nested.getAggregations().asList().isEmpty(), is(false));
+
+                Terms usernames = nested.getAggregations().get("field2");
+                assertThat(usernames, notNullValue());
+                assertThat(usernames.getBuckets().size(), equalTo(5));
+                List<Terms.Bucket> usernameBuckets = new ArrayList<>(usernames.getBuckets());
+
+                Terms.Bucket bucket = usernameBuckets.get(0);
+                assertThat(bucket.getKeyAsString(), equalTo("0"));
+                assertThat(bucket.getDocCount(), equalTo(12L));
+                ReverseNested reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                assertThat(reverseNested.getDocCount(), equalTo(5L));
+                Terms tags = reverseNested.getAggregations().get("field1");
+                List<Terms.Bucket> tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(2));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(3L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(2L));
+
+                bucket = usernameBuckets.get(1);
+                assertThat(bucket.getKeyAsString(), equalTo("1"));
+                assertThat(bucket.getDocCount(), equalTo(6L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                assertThat(reverseNested.getDocCount(), equalTo(4L));
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(4));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+
+                bucket = usernameBuckets.get(2);
+                assertThat(bucket.getKeyAsString(), equalTo("2"));
+                assertThat(bucket.getDocCount(), equalTo(5L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                assertThat(reverseNested.getDocCount(), equalTo(4L));
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(4));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("a"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("b"));
+                assertThat(tagsBuckets.get(1).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(2).getKeyAsString(), equalTo("c"));
+                assertThat(tagsBuckets.get(2).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(3).getKeyAsString(), equalTo("e"));
+                assertThat(tagsBuckets.get(3).getDocCount(), equalTo(1L));
+
+                bucket = usernameBuckets.get(3);
+                assertThat(bucket.getKeyAsString(), equalTo("3"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                assertThat(reverseNested.getDocCount(), equalTo(2L));
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(2));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("f"));
+
+                bucket = usernameBuckets.get(4);
+                assertThat(bucket.getKeyAsString(), equalTo("4"));
+                assertThat(bucket.getDocCount(), equalTo(2L));
+                reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                assertThat(reverseNested.getDocCount(), equalTo(2L));
+                tags = reverseNested.getAggregations().get("field1");
+                tagsBuckets = new ArrayList<>(tags.getBuckets());
+                assertThat(tagsBuckets.size(), equalTo(2));
+                assertThat(tagsBuckets.get(0).getKeyAsString(), equalTo("d"));
+                assertThat(tagsBuckets.get(0).getDocCount(), equalTo(1L));
+                assertThat(tagsBuckets.get(1).getKeyAsString(), equalTo("f"));
+            }
+        );
     }
 
     public void testReverseNestedAggWithoutNestedAgg() {
@@ -467,26 +470,32 @@ public class ReverseNestedIT extends ESIntegTestCase {
     }
 
     public void testNonExistingNestedField() throws Exception {
-        SearchResponse searchResponse = prepareSearch("idx2").setQuery(matchAllQuery())
-            .addAggregation(nested("nested2", "nested1.nested2").subAggregation(reverseNested("incorrect").path("nested3")))
-            .get();
+        assertNoFailuresAndResponse(
+            prepareSearch("idx2").setQuery(matchAllQuery())
+                .addAggregation(nested("nested2", "nested1.nested2").subAggregation(reverseNested("incorrect").path("nested3"))),
+            response -> {
 
-        Nested nested = searchResponse.getAggregations().get("nested2");
-        assertThat(nested, notNullValue());
-        assertThat(nested.getName(), equalTo("nested2"));
+                Nested nested = response.getAggregations().get("nested2");
+                assertThat(nested, notNullValue());
+                assertThat(nested.getName(), equalTo("nested2"));
 
-        ReverseNested reverseNested = nested.getAggregations().get("incorrect");
-        assertThat(reverseNested.getDocCount(), is(0L));
+                ReverseNested reverseNested = nested.getAggregations().get("incorrect");
+                assertThat(reverseNested.getDocCount(), is(0L));
+            }
+        );
 
         // Test that parsing the reverse_nested agg doesn't fail, because the parent nested agg is unmapped:
-        searchResponse = prepareSearch("idx1").setQuery(matchAllQuery())
-            .addAggregation(nested("incorrect1", "incorrect1").subAggregation(reverseNested("incorrect2").path("incorrect2")))
-            .get();
-
-        nested = searchResponse.getAggregations().get("incorrect1");
-        assertThat(nested, notNullValue());
-        assertThat(nested.getName(), equalTo("incorrect1"));
-        assertThat(nested.getDocCount(), is(0L));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx1").setQuery(matchAllQuery())
+                .addAggregation(nested("incorrect1", "incorrect1").subAggregation(reverseNested("incorrect2").path("incorrect2"))),
+            response -> {
+
+                Nested nested = response.getAggregations().get("incorrect1");
+                assertThat(nested, notNullValue());
+                assertThat(nested.getName(), equalTo("incorrect1"));
+                assertThat(nested.getDocCount(), is(0L));
+            }
+        );
     }
 
     public void testSameParentDocHavingMultipleBuckets() throws Exception {
@@ -603,110 +612,117 @@ public class ReverseNestedIT extends ESIntegTestCase {
             )
             .get();
 
-        SearchResponse response = prepareSearch("idx3").addAggregation(
-            nested("nested_0", "category").subAggregation(
-                terms("group_by_category").field("category.name")
-                    .subAggregation(
-                        reverseNested("to_root").subAggregation(
-                            nested("nested_1", "sku").subAggregation(
-                                filter("filter_by_sku", termQuery("sku.sku_type", "bar1")).subAggregation(
-                                    count("sku_count").field("sku.sku_type")
+        assertNoFailuresAndResponse(
+            prepareSearch("idx3").addAggregation(
+                nested("nested_0", "category").subAggregation(
+                    terms("group_by_category").field("category.name")
+                        .subAggregation(
+                            reverseNested("to_root").subAggregation(
+                                nested("nested_1", "sku").subAggregation(
+                                    filter("filter_by_sku", termQuery("sku.sku_type", "bar1")).subAggregation(
+                                        count("sku_count").field("sku.sku_type")
+                                    )
                                 )
                             )
                         )
-                    )
-            )
-        ).get();
-        assertNoFailures(response);
-        assertHitCount(response, 1);
-
-        Nested nested0 = response.getAggregations().get("nested_0");
-        assertThat(nested0.getDocCount(), equalTo(3L));
-        Terms terms = nested0.getAggregations().get("group_by_category");
-        assertThat(terms.getBuckets().size(), equalTo(3));
-        for (String bucketName : new String[] { "abc", "klm", "xyz" }) {
-            logger.info("Checking results for bucket {}", bucketName);
-            Terms.Bucket bucket = terms.getBucketByKey(bucketName);
-            assertThat(bucket.getDocCount(), equalTo(1L));
-            ReverseNested toRoot = bucket.getAggregations().get("to_root");
-            assertThat(toRoot.getDocCount(), equalTo(1L));
-            Nested nested1 = toRoot.getAggregations().get("nested_1");
-            assertThat(nested1.getDocCount(), equalTo(5L));
-            Filter filterByBar = nested1.getAggregations().get("filter_by_sku");
-            assertThat(filterByBar.getDocCount(), equalTo(3L));
-            ValueCount barCount = filterByBar.getAggregations().get("sku_count");
-            assertThat(barCount.getValue(), equalTo(3L));
-        }
+                )
+            ),
+            response -> {
+                assertHitCount(response, 1);
+
+                Nested nested0 = response.getAggregations().get("nested_0");
+                assertThat(nested0.getDocCount(), equalTo(3L));
+                Terms terms = nested0.getAggregations().get("group_by_category");
+                assertThat(terms.getBuckets().size(), equalTo(3));
+                for (String bucketName : new String[] { "abc", "klm", "xyz" }) {
+                    logger.info("Checking results for bucket {}", bucketName);
+                    Terms.Bucket bucket = terms.getBucketByKey(bucketName);
+                    assertThat(bucket.getDocCount(), equalTo(1L));
+                    ReverseNested toRoot = bucket.getAggregations().get("to_root");
+                    assertThat(toRoot.getDocCount(), equalTo(1L));
+                    Nested nested1 = toRoot.getAggregations().get("nested_1");
+                    assertThat(nested1.getDocCount(), equalTo(5L));
+                    Filter filterByBar = nested1.getAggregations().get("filter_by_sku");
+                    assertThat(filterByBar.getDocCount(), equalTo(3L));
+                    ValueCount barCount = filterByBar.getAggregations().get("sku_count");
+                    assertThat(barCount.getValue(), equalTo(3L));
+                }
+            }
+        );
 
-        response = prepareSearch("idx3").addAggregation(
-            nested("nested_0", "category").subAggregation(
-                terms("group_by_category").field("category.name")
-                    .subAggregation(
-                        reverseNested("to_root").subAggregation(
-                            nested("nested_1", "sku").subAggregation(
-                                filter("filter_by_sku", termQuery("sku.sku_type", "bar1")).subAggregation(
-                                    nested("nested_2", "sku.colors").subAggregation(
-                                        filter("filter_sku_color", termQuery("sku.colors.name", "red")).subAggregation(
-                                            reverseNested("reverse_to_sku").path("sku")
-                                                .subAggregation(count("sku_count").field("sku.sku_type"))
+        assertNoFailuresAndResponse(
+            prepareSearch("idx3").addAggregation(
+                nested("nested_0", "category").subAggregation(
+                    terms("group_by_category").field("category.name")
+                        .subAggregation(
+                            reverseNested("to_root").subAggregation(
+                                nested("nested_1", "sku").subAggregation(
+                                    filter("filter_by_sku", termQuery("sku.sku_type", "bar1")).subAggregation(
+                                        nested("nested_2", "sku.colors").subAggregation(
+                                            filter("filter_sku_color", termQuery("sku.colors.name", "red")).subAggregation(
+                                                reverseNested("reverse_to_sku").path("sku")
+                                                    .subAggregation(count("sku_count").field("sku.sku_type"))
+                                            )
                                         )
                                     )
                                 )
                             )
                         )
-                    )
-            )
-        ).get();
-        assertNoFailures(response);
-        assertHitCount(response, 1);
-
-        nested0 = response.getAggregations().get("nested_0");
-        assertThat(nested0.getDocCount(), equalTo(3L));
-        terms = nested0.getAggregations().get("group_by_category");
-        assertThat(terms.getBuckets().size(), equalTo(3));
-        for (String bucketName : new String[] { "abc", "klm", "xyz" }) {
-            logger.info("Checking results for bucket {}", bucketName);
-            Terms.Bucket bucket = terms.getBucketByKey(bucketName);
-            assertThat(bucket.getDocCount(), equalTo(1L));
-            ReverseNested toRoot = bucket.getAggregations().get("to_root");
-            assertThat(toRoot.getDocCount(), equalTo(1L));
-            Nested nested1 = toRoot.getAggregations().get("nested_1");
-            assertThat(nested1.getDocCount(), equalTo(5L));
-            Filter filterByBar = nested1.getAggregations().get("filter_by_sku");
-            assertThat(filterByBar.getDocCount(), equalTo(3L));
-            Nested nested2 = filterByBar.getAggregations().get("nested_2");
-            assertThat(nested2.getDocCount(), equalTo(8L));
-            Filter filterBarColor = nested2.getAggregations().get("filter_sku_color");
-            assertThat(filterBarColor.getDocCount(), equalTo(2L));
-            ReverseNested reverseToBar = filterBarColor.getAggregations().get("reverse_to_sku");
-            assertThat(reverseToBar.getDocCount(), equalTo(2L));
-            ValueCount barCount = reverseToBar.getAggregations().get("sku_count");
-            assertThat(barCount.getValue(), equalTo(2L));
-        }
+                )
+            ),
+            response -> {
+                assertHitCount(response, 1);
+
+                Nested nested0 = response.getAggregations().get("nested_0");
+                assertThat(nested0.getDocCount(), equalTo(3L));
+                Terms terms = nested0.getAggregations().get("group_by_category");
+                assertThat(terms.getBuckets().size(), equalTo(3));
+                for (String bucketName : new String[] { "abc", "klm", "xyz" }) {
+                    logger.info("Checking results for bucket {}", bucketName);
+                    Terms.Bucket bucket = terms.getBucketByKey(bucketName);
+                    assertThat(bucket.getDocCount(), equalTo(1L));
+                    ReverseNested toRoot = bucket.getAggregations().get("to_root");
+                    assertThat(toRoot.getDocCount(), equalTo(1L));
+                    Nested nested1 = toRoot.getAggregations().get("nested_1");
+                    assertThat(nested1.getDocCount(), equalTo(5L));
+                    Filter filterByBar = nested1.getAggregations().get("filter_by_sku");
+                    assertThat(filterByBar.getDocCount(), equalTo(3L));
+                    Nested nested2 = filterByBar.getAggregations().get("nested_2");
+                    assertThat(nested2.getDocCount(), equalTo(8L));
+                    Filter filterBarColor = nested2.getAggregations().get("filter_sku_color");
+                    assertThat(filterBarColor.getDocCount(), equalTo(2L));
+                    ReverseNested reverseToBar = filterBarColor.getAggregations().get("reverse_to_sku");
+                    assertThat(reverseToBar.getDocCount(), equalTo(2L));
+                    ValueCount barCount = reverseToBar.getAggregations().get("sku_count");
+                    assertThat(barCount.getValue(), equalTo(2L));
+                }
+            }
+        );
     }
 
     public void testFieldAlias() {
-        SearchResponse response = prepareSearch("idx1").addAggregation(
-            nested("nested1", "nested1").subAggregation(
-                terms("field2").field("nested1.field2")
-                    .subAggregation(
-                        reverseNested("nested1_to_field1").subAggregation(
-                            terms("field1").field("alias").collectMode(randomFrom(SubAggCollectionMode.values()))
+        assertNoFailuresAndResponse(
+            prepareSearch("idx1").addAggregation(
+                nested("nested1", "nested1").subAggregation(
+                    terms("field2").field("nested1.field2")
+                        .subAggregation(
+                            reverseNested("nested1_to_field1").subAggregation(
+                                terms("field1").field("alias").collectMode(randomFrom(SubAggCollectionMode.values()))
+                            )
                         )
-                    )
-            )
-        ).get();
-
-        assertNoFailures(response);
-
-        Nested nested = response.getAggregations().get("nested1");
-        Terms nestedTerms = nested.getAggregations().get("field2");
-        Terms.Bucket bucket = nestedTerms.getBuckets().iterator().next();
-
-        ReverseNested reverseNested = bucket.getAggregations().get("nested1_to_field1");
-        Terms reverseNestedTerms = reverseNested.getAggregations().get("field1");
-
-        assertThat(((InternalAggregation) reverseNested).getProperty("field1"), sameInstance(reverseNestedTerms));
-        assertThat(reverseNestedTerms.getBuckets().size(), equalTo(6));
+                )
+            ),
+            response -> {
+                Nested nested = response.getAggregations().get("nested1");
+                Terms nestedTerms = nested.getAggregations().get("field2");
+                Terms.Bucket bucket = nestedTerms.getBuckets().iterator().next();
+
+                ReverseNested reverseNested = bucket.getAggregations().get("nested1_to_field1");
+                Terms reverseNestedTerms = reverseNested.getAggregations().get("field1");
+
+                assertThat(((InternalAggregation) reverseNested).getProperty("field1"), sameInstance(reverseNestedTerms));
+                assertThat(reverseNestedTerms.getBuckets().size(), equalTo(6));
+            }
+        );
     }
 }

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/SamplerIT.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.elasticsearch.action.admin.indices.refresh.RefreshRequest;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.SearchType;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.search.aggregations.BucketOrder;
@@ -27,6 +26,7 @@ import static org.elasticsearch.search.aggregations.AggregationBuilders.sampler;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
@@ -91,96 +91,103 @@ public class SamplerIT extends ESIntegTestCase {
         // Tests that we can refer to nested elements under a sample in a path
         // statement
         boolean asc = randomBoolean();
-        SearchResponse response = prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH)
-            .addAggregation(
-                terms("genres").field("genre")
-                    .order(BucketOrder.aggregation("sample>max_price.value", asc))
-                    .subAggregation(sampler("sample").shardSize(100).subAggregation(max("max_price").field("price")))
-            )
-            .get();
-        assertNoFailures(response);
-        Terms genres = response.getAggregations().get("genres");
-        List<? extends Bucket> genreBuckets = genres.getBuckets();
-        // For this test to be useful we need >1 genre bucket to compare
-        assertThat(genreBuckets.size(), greaterThan(1));
-        double lastMaxPrice = asc ? Double.MIN_VALUE : Double.MAX_VALUE;
-        for (Terms.Bucket genreBucket : genres.getBuckets()) {
-            Sampler sample = genreBucket.getAggregations().get("sample");
-            Max maxPriceInGenre = sample.getAggregations().get("max_price");
-            double price = maxPriceInGenre.value();
-            if (asc) {
-                assertThat(price, greaterThanOrEqualTo(lastMaxPrice));
-            } else {
-                assertThat(price, lessThanOrEqualTo(lastMaxPrice));
+        assertNoFailuresAndResponse(
+            prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH)
+                .addAggregation(
+                    terms("genres").field("genre")
+                        .order(BucketOrder.aggregation("sample>max_price.value", asc))
+                        .subAggregation(sampler("sample").shardSize(100).subAggregation(max("max_price").field("price")))
+                ),
+            response -> {
+                Terms genres = response.getAggregations().get("genres");
+                List<? extends Bucket> genreBuckets = genres.getBuckets();
+                // For this test to be useful we need >1 genre bucket to compare
+                assertThat(genreBuckets.size(), greaterThan(1));
+                double lastMaxPrice = asc ? Double.MIN_VALUE : Double.MAX_VALUE;
+                for (Terms.Bucket genreBucket : genres.getBuckets()) {
+                    Sampler sample = genreBucket.getAggregations().get("sample");
+                    Max maxPriceInGenre = sample.getAggregations().get("max_price");
+                    double price = maxPriceInGenre.value();
+                    if (asc) {
+                        assertThat(price, greaterThanOrEqualTo(lastMaxPrice));
+                    } else {
+                        assertThat(price, lessThanOrEqualTo(lastMaxPrice));
+                    }
+                    lastMaxPrice = price;
+                }
             }
-            lastMaxPrice = price;
-        }
-
+        );
     }
 
     public void testSimpleSampler() throws Exception {
         SamplerAggregationBuilder sampleAgg = sampler("sample").shardSize(100);
         sampleAgg.subAggregation(terms("authors").field("author"));
-        SearchResponse response = prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH)
-            .setQuery(new TermQueryBuilder("genre", "fantasy"))
-            .setFrom(0)
-            .setSize(60)
-            .addAggregation(sampleAgg)
-            .get();
-        assertNoFailures(response);
-        Sampler sample = response.getAggregations().get("sample");
-        Terms authors = sample.getAggregations().get("authors");
-        List<? extends Bucket> testBuckets = authors.getBuckets();
-
-        long maxBooksPerAuthor = 0;
-        for (Terms.Bucket testBucket : testBuckets) {
-            maxBooksPerAuthor = Math.max(testBucket.getDocCount(), maxBooksPerAuthor);
-        }
-        assertThat(maxBooksPerAuthor, equalTo(3L));
+        assertNoFailuresAndResponse(
+            prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH)
+                .setQuery(new TermQueryBuilder("genre", "fantasy"))
+                .setFrom(0)
+                .setSize(60)
+                .addAggregation(sampleAgg),
+            response -> {
+                Sampler sample = response.getAggregations().get("sample");
+                Terms authors = sample.getAggregations().get("authors");
+                List<? extends Bucket> testBuckets = authors.getBuckets();
+
+                long maxBooksPerAuthor = 0;
+                for (Terms.Bucket testBucket : testBuckets) {
+                    maxBooksPerAuthor = Math.max(testBucket.getDocCount(), maxBooksPerAuthor);
+                }
+                assertThat(maxBooksPerAuthor, equalTo(3L));
+            }
+        );
     }
 
     public void testUnmappedChildAggNoDiversity() throws Exception {
         SamplerAggregationBuilder sampleAgg = sampler("sample").shardSize(100);
         sampleAgg.subAggregation(terms("authors").field("author"));
-        SearchResponse response = prepareSearch("idx_unmapped").setSearchType(SearchType.QUERY_THEN_FETCH)
-            .setQuery(new TermQueryBuilder("genre", "fantasy"))
-            .setFrom(0)
-            .setSize(60)
-            .addAggregation(sampleAgg)
-            .get();
-        assertNoFailures(response);
-        Sampler sample = response.getAggregations().get("sample");
-        assertThat(sample.getDocCount(), equalTo(0L));
-        Terms authors = sample.getAggregations().get("authors");
-        assertThat(authors.getBuckets().size(), equalTo(0));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx_unmapped").setSearchType(SearchType.QUERY_THEN_FETCH)
+                .setQuery(new TermQueryBuilder("genre", "fantasy"))
+                .setFrom(0)
+                .setSize(60)
+                .addAggregation(sampleAgg),
+            response -> {
+                Sampler sample = response.getAggregations().get("sample");
+                assertThat(sample.getDocCount(), equalTo(0L));
+                Terms authors = sample.getAggregations().get("authors");
+                assertThat(authors.getBuckets().size(), equalTo(0));
+            }
+        );
     }
 
     public void testPartiallyUnmappedChildAggNoDiversity() throws Exception {
         SamplerAggregationBuilder sampleAgg = sampler("sample").shardSize(100);
         sampleAgg.subAggregation(terms("authors").field("author"));
-        SearchResponse response = prepareSearch("idx_unmapped", "test").setSearchType(SearchType.QUERY_THEN_FETCH)
-            .setQuery(new TermQueryBuilder("genre", "fantasy"))
-            .setFrom(0)
-            .setSize(60)
-            .setExplain(true)
-            .addAggregation(sampleAgg)
-            .get();
-        assertNoFailures(response);
-        Sampler sample = response.getAggregations().get("sample");
-        assertThat(sample.getDocCount(), greaterThan(0L));
-        Terms authors = sample.getAggregations().get("authors");
-        assertThat(authors.getBuckets().size(), greaterThan(0));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx_unmapped", "test").setSearchType(SearchType.QUERY_THEN_FETCH)
+                .setQuery(new TermQueryBuilder("genre", "fantasy"))
+                .setFrom(0)
+                .setSize(60)
+                .setExplain(true)
+                .addAggregation(sampleAgg),
+            response -> {
+                Sampler sample = response.getAggregations().get("sample");
+                assertThat(sample.getDocCount(), greaterThan(0L));
+                Terms authors = sample.getAggregations().get("authors");
+                assertThat(authors.getBuckets().size(), greaterThan(0));
+            }
+        );
     }
 
     public void testRidiculousShardSizeSampler() throws Exception {
         SamplerAggregationBuilder sampleAgg = sampler("sample").shardSize(Integer.MAX_VALUE);
         sampleAgg.subAggregation(terms("authors").field("author"));
-        SearchResponse response = prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH)
-            .setQuery(new TermQueryBuilder("genre", "fantasy"))
-            .setFrom(0)
-            .setSize(60)
-            .addAggregation(sampleAgg)
-            .get();
-        assertNoFailures(response);
+        assertNoFailures(
+            prepareSearch("test").setSearchType(SearchType.QUERY_THEN_FETCH)
+                .setQuery(new TermQueryBuilder("genre", "fantasy"))
+                .setFrom(0)
+                .setSize(60)
+                .addAggregation(sampleAgg)
+        );
     }
 }

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/ShardReduceIT.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.elasticsearch.action.index.IndexRequestBuilder;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.geometry.utils.Geohash;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.search.aggregations.Aggregator.SubAggCollectionMode;
@@ -36,7 +35,7 @@ import static org.elasticsearch.search.aggregations.AggregationBuilders.nested;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.range;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
-import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.hamcrest.Matchers.equalTo;
 
@@ -87,246 +86,248 @@ public class ShardReduceIT extends ESIntegTestCase {
     }
 
     public void testGlobal() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                global("global").subAggregation(
-                    dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
-                )
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Global global = response.getAggregations().get("global");
-        Histogram histo = global.getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    global("global").subAggregation(
+                        dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
+                    )
+                ),
+            response -> {
+                Global global = response.getAggregations().get("global");
+                Histogram histo = global.getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testFilter() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                filter("filter", QueryBuilders.matchAllQuery()).subAggregation(
-                    dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
-                )
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Filter filter = response.getAggregations().get("filter");
-        Histogram histo = filter.getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    filter("filter", QueryBuilders.matchAllQuery()).subAggregation(
+                        dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
+                    )
+                ),
+            response -> {
+                Filter filter = response.getAggregations().get("filter");
+                Histogram histo = filter.getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testMissing() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                missing("missing").field("foobar")
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Missing missing = response.getAggregations().get("missing");
-        Histogram histo = missing.getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    missing("missing").field("foobar")
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Missing missing = response.getAggregations().get("missing");
+                Histogram histo = missing.getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testGlobalWithFilterWithMissing() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                global("global").subAggregation(
-                    filter("filter", QueryBuilders.matchAllQuery()).subAggregation(
-                        missing("missing").field("foobar")
-                            .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    global("global").subAggregation(
+                        filter("filter", QueryBuilders.matchAllQuery()).subAggregation(
+                            missing("missing").field("foobar")
+                                .subAggregation(
+                                    dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
+                                )
+                        )
                     )
-                )
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Global global = response.getAggregations().get("global");
-        Filter filter = global.getAggregations().get("filter");
-        Missing missing = filter.getAggregations().get("missing");
-        Histogram histo = missing.getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+                ),
+            response -> {
+                Global global = response.getAggregations().get("global");
+                Filter filter = global.getAggregations().get("filter");
+                Missing missing = filter.getAggregations().get("missing");
+                Histogram histo = missing.getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testNested() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                nested("nested", "nested").subAggregation(
-                    dateHistogram("histo").field("nested.date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
-                )
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Nested nested = response.getAggregations().get("nested");
-        Histogram histo = nested.getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    nested("nested", "nested").subAggregation(
+                        dateHistogram("histo").field("nested.date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0)
+                    )
+                ),
+            response -> {
+                Nested nested = response.getAggregations().get("nested");
+                Histogram histo = nested.getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testStringTerms() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                terms("terms").field("term-s")
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Terms terms = response.getAggregations().get("terms");
-        Histogram histo = terms.getBucketByKey("term").getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    terms("terms").field("term-s")
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("terms");
+                Histogram histo = terms.getBucketByKey("term").getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testLongTerms() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                terms("terms").field("term-l")
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Terms terms = response.getAggregations().get("terms");
-        Histogram histo = terms.getBucketByKey("1").getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    terms("terms").field("term-l")
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("terms");
+                Histogram histo = terms.getBucketByKey("1").getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testDoubleTerms() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                terms("terms").field("term-d")
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Terms terms = response.getAggregations().get("terms");
-        Histogram histo = terms.getBucketByKey("1.5").getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    terms("terms").field("term-d")
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("terms");
+                Histogram histo = terms.getBucketByKey("1.5").getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testRange() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                range("range").field("value")
-                    .addRange("r1", 0, 10)
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Range range = response.getAggregations().get("range");
-        Histogram histo = range.getBuckets().get(0).getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    range("range").field("value")
+                        .addRange("r1", 0, 10)
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Range range = response.getAggregations().get("range");
+                Histogram histo = range.getBuckets().get(0).getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testDateRange() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                dateRange("range").field("date")
-                    .addRange("r1", "2014-01-01", "2014-01-10")
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Range range = response.getAggregations().get("range");
-        Histogram histo = range.getBuckets().get(0).getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    dateRange("range").field("date")
+                        .addRange("r1", "2014-01-01", "2014-01-10")
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Range range = response.getAggregations().get("range");
+                Histogram histo = range.getBuckets().get(0).getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testIpRange() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                ipRange("range").field("ip")
-                    .addRange("r1", "10.0.0.1", "10.0.0.10")
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Range range = response.getAggregations().get("range");
-        Histogram histo = range.getBuckets().get(0).getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    ipRange("range").field("ip")
+                        .addRange("r1", "10.0.0.1", "10.0.0.10")
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Range range = response.getAggregations().get("range");
+                Histogram histo = range.getBuckets().get(0).getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testHistogram() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                histogram("topHisto").field("value")
-                    .interval(5)
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Histogram topHisto = response.getAggregations().get("topHisto");
-        Histogram histo = topHisto.getBuckets().get(0).getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    histogram("topHisto").field("value")
+                        .interval(5)
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Histogram topHisto = response.getAggregations().get("topHisto");
+                Histogram histo = topHisto.getBuckets().get(0).getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testDateHistogram() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                dateHistogram("topHisto").field("date")
-                    .calendarInterval(DateHistogramInterval.MONTH)
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        Histogram topHisto = response.getAggregations().get("topHisto");
-        Histogram histo = topHisto.getBuckets().iterator().next().getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    dateHistogram("topHisto").field("date")
+                        .calendarInterval(DateHistogramInterval.MONTH)
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                Histogram topHisto = response.getAggregations().get("topHisto");
+                Histogram histo = topHisto.getBuckets().iterator().next().getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
 
     }
 
     public void testGeoHashGrid() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                geohashGrid("grid").field("location")
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        GeoGrid grid = response.getAggregations().get("grid");
-        Histogram histo = grid.getBuckets().iterator().next().getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    geohashGrid("grid").field("location")
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                GeoGrid grid = response.getAggregations().get("grid");
+                Histogram histo = grid.getBuckets().iterator().next().getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
     public void testGeoTileGrid() throws Exception {
-        SearchResponse response = prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
-            .addAggregation(
-                geotileGrid("grid").field("location")
-                    .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
-            )
-            .get();
-
-        assertNoFailures(response);
-
-        GeoGrid grid = response.getAggregations().get("grid");
-        Histogram histo = grid.getBuckets().iterator().next().getAggregations().get("histo");
-        assertThat(histo.getBuckets().size(), equalTo(4));
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(QueryBuilders.matchAllQuery())
+                .addAggregation(
+                    geotileGrid("grid").field("location")
+                        .subAggregation(dateHistogram("histo").field("date").fixedInterval(DateHistogramInterval.DAY).minDocCount(0))
+                ),
+            response -> {
+                GeoGrid grid = response.getAggregations().get("grid");
+                Histogram histo = grid.getBuckets().iterator().next().getAggregations().get("histo");
+                assertThat(histo.getBuckets().size(), equalTo(4));
+            }
+        );
     }
 
 }

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/ShardSizeTermsIT.java

@@ -7,7 +7,6 @@
  */
 package org.elasticsearch.search.aggregations.bucket;
 
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.search.aggregations.Aggregator.SubAggCollectionMode;
 import org.elasticsearch.search.aggregations.BucketOrder;
 import org.elasticsearch.search.aggregations.bucket.terms.Terms;
@@ -18,6 +17,7 @@ import java.util.Map;
 
 import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery;
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.hamcrest.Matchers.equalTo;
 
 public class ShardSizeTermsIT extends ShardSizeTestCase {
@@ -26,22 +26,27 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<String, Long> expected = new HashMap<>();
-        expected.put("1", 8L);
-        expected.put("3", 8L);
-        expected.put("2", 5L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<String, Long> expected = new HashMap<>();
+                expected.put("1", 8L);
+                expected.put("3", 8L);
+                expected.put("2", 5L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
+                }
+            }
+        );
     }
 
     public void testShardSizeEqualsSizeString() throws Exception {
@@ -49,26 +54,28 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .shardSize(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<String, Long> expected = new HashMap<>();
-        expected.put("1", 8L);
-        expected.put("3", 8L);
-        expected.put("2", 4L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .shardSize(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<String, Long> expected = new HashMap<>();
+                expected.put("1", 8L);
+                expected.put("3", 8L);
+                expected.put("2", 4L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
+                }
+            }
+        );
     }
 
     public void testWithShardSizeString() throws Exception {
@@ -77,26 +84,28 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .shardSize(5)
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
-        Map<String, Long> expected = new HashMap<>();
-        expected.put("1", 8L);
-        expected.put("3", 8L);
-        expected.put("2", 5L); // <-- count is now fixed
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .shardSize(5)
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
+                Map<String, Long> expected = new HashMap<>();
+                expected.put("1", 8L);
+                expected.put("3", 8L);
+                expected.put("2", 5L); // <-- count is now fixed
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
+                }
+            }
+        );
     }
 
     public void testWithShardSizeStringSingleShard() throws Exception {
@@ -105,27 +114,29 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setRouting(routing1)
-            .setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .shardSize(5)
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
-        Map<String, Long> expected = new HashMap<>();
-        expected.put("1", 5L);
-        expected.put("2", 4L);
-        expected.put("3", 3L); // <-- count is now fixed
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKey())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setRouting(routing1)
+                .setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .shardSize(5)
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
+                Map<String, Long> expected = new HashMap<>();
+                expected.put("1", 5L);
+                expected.put("2", 4L);
+                expected.put("3", 3L); // <-- count is now fixed
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKey())));
+                }
+            }
+        );
     }
 
     public void testNoShardSizeTermOrderString() throws Exception {
@@ -133,22 +144,24 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.key(true))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<String, Long> expected = new HashMap<>();
-        expected.put("1", 8L);
-        expected.put("2", 5L);
-        expected.put("3", 8L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.key(true))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<String, Long> expected = new HashMap<>();
+                expected.put("1", 8L);
+                expected.put("2", 5L);
+                expected.put("3", 8L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsString())));
+                }
+            }
+        );
     }
 
     public void testNoShardSizeLong() throws Exception {
@@ -156,22 +169,27 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(3, 8L);
-        expected.put(2, 5L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(3, 8L);
+                expected.put(2, 5L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testShardSizeEqualsSizeLong() throws Exception {
@@ -179,26 +197,28 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .shardSize(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(3, 8L);
-        expected.put(2, 4L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .shardSize(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(3, 8L);
+                expected.put(2, 4L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testWithShardSizeLong() throws Exception {
@@ -206,26 +226,28 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .shardSize(5)
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(3, 8L);
-        expected.put(2, 5L); // <-- count is now fixed
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .shardSize(5)
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(3, 8L);
+                expected.put(2, 5L); // <-- count is now fixed
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testWithShardSizeLongSingleShard() throws Exception {
@@ -234,27 +256,29 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setRouting(routing1)
-            .setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .shardSize(5)
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 5L);
-        expected.put(2, 4L);
-        expected.put(3, 3L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setRouting(routing1)
+                .setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .shardSize(5)
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3)); // we still only return 3 entries (based on the 'size' param)
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 5L);
+                expected.put(2, 4L);
+                expected.put(3, 3L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testNoShardSizeTermOrderLong() throws Exception {
@@ -262,22 +286,24 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.key(true))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(2, 5L);
-        expected.put(3, 8L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.key(true))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(2, 5L);
+                expected.put(3, 8L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testNoShardSizeDouble() throws Exception {
@@ -285,22 +311,27 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(3, 8L);
-        expected.put(2, 5L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(3, 8L);
+                expected.put(2, 5L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testShardSizeEqualsSizeDouble() throws Exception {
@@ -308,26 +339,28 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .shardSize(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(3, 8L);
-        expected.put(2, 4L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .shardSize(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(3, 8L);
+                expected.put(2, 4L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testWithShardSizeDouble() throws Exception {
@@ -335,26 +368,28 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .shardSize(5)
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(3, 8L);
-        expected.put(2, 5L); // <-- count is now fixed
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .shardSize(5)
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(3, 8L);
+                expected.put(2, 5L); // <-- count is now fixed
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testWithShardSizeDoubleSingleShard() throws Exception {
@@ -362,27 +397,29 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setRouting(routing1)
-            .setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key")
-                    .size(3)
-                    .collectMode(randomFrom(SubAggCollectionMode.values()))
-                    .shardSize(5)
-                    .order(BucketOrder.count(false))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 5L);
-        expected.put(2, 4L);
-        expected.put(3, 3L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setRouting(routing1)
+                .setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key")
+                        .size(3)
+                        .collectMode(randomFrom(SubAggCollectionMode.values()))
+                        .shardSize(5)
+                        .order(BucketOrder.count(false))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 5L);
+                expected.put(2, 4L);
+                expected.put(3, 3L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 
     public void testNoShardSizeTermOrderDouble() throws Exception {
@@ -390,21 +427,23 @@ public class ShardSizeTermsIT extends ShardSizeTestCase {
 
         indexData();
 
-        SearchResponse response = prepareSearch("idx").setQuery(matchAllQuery())
-            .addAggregation(
-                terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.key(true))
-            )
-            .get();
-
-        Terms terms = response.getAggregations().get("keys");
-        List<? extends Terms.Bucket> buckets = terms.getBuckets();
-        assertThat(buckets.size(), equalTo(3));
-        Map<Integer, Long> expected = new HashMap<>();
-        expected.put(1, 8L);
-        expected.put(2, 5L);
-        expected.put(3, 8L);
-        for (Terms.Bucket bucket : buckets) {
-            assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
-        }
+        assertNoFailuresAndResponse(
+            prepareSearch("idx").setQuery(matchAllQuery())
+                .addAggregation(
+                    terms("keys").field("key").size(3).collectMode(randomFrom(SubAggCollectionMode.values())).order(BucketOrder.key(true))
+                ),
+            response -> {
+                Terms terms = response.getAggregations().get("keys");
+                List<? extends Terms.Bucket> buckets = terms.getBuckets();
+                assertThat(buckets.size(), equalTo(3));
+                Map<Integer, Long> expected = new HashMap<>();
+                expected.put(1, 8L);
+                expected.put(2, 5L);
+                expected.put(3, 8L);
+                for (Terms.Bucket bucket : buckets) {
+                    assertThat(bucket.getDocCount(), equalTo(expected.get(bucket.getKeyAsNumber().intValue())));
+                }
+            }
+        );
     }
 }

server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/bucket/TermsShardMinDocCountIT.java

@@ -8,7 +8,6 @@
 package org.elasticsearch.search.aggregations.bucket;
 
 import org.elasticsearch.action.index.IndexRequestBuilder;
-import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.search.aggregations.BucketOrder;
 import org.elasticsearch.search.aggregations.bucket.filter.InternalFilter;
@@ -26,6 +25,7 @@ import static org.elasticsearch.search.aggregations.AggregationBuilders.signific
 import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailuresAndResponse;
 import static org.hamcrest.Matchers.equalTo;
 
 public class TermsShardMinDocCountIT extends ESIntegTestCase {
@@ -61,34 +61,41 @@ public class TermsShardMinDocCountIT extends ESIntegTestCase {
         indexRandom(true, false, indexBuilders);
 
         // first, check that indeed when not setting the shardMinDocCount parameter 0 terms are returned
-        SearchResponse response = prepareSearch(index).addAggregation(
-            (filter("inclass", QueryBuilders.termQuery("class", true))).subAggregation(
-                significantTerms("mySignificantTerms").field("text")
-                    .minDocCount(2)
-                    .size(2)
-                    .shardSize(2)
-                    .executionHint(randomExecutionHint())
-            )
-        ).get();
-        assertNoFailures(response);
-        InternalFilter filteredBucket = response.getAggregations().get("inclass");
-        SignificantTerms sigterms = filteredBucket.getAggregations().get("mySignificantTerms");
-        assertThat(sigterms.getBuckets().size(), equalTo(0));
-
-        response = prepareSearch(index).addAggregation(
-            (filter("inclass", QueryBuilders.termQuery("class", true))).subAggregation(
-                significantTerms("mySignificantTerms").field("text")
-                    .minDocCount(2)
-                    .shardSize(2)
-                    .shardMinDocCount(2)
-                    .size(2)
-                    .executionHint(randomExecutionHint())
-            )
-        ).get();
-        assertNoFailures(response);
-        filteredBucket = response.getAggregations().get("inclass");
-        sigterms = filteredBucket.getAggregations().get("mySignificantTerms");
-        assertThat(sigterms.getBuckets().size(), equalTo(2));
+        assertNoFailuresAndResponse(
+            prepareSearch(index).addAggregation(
+                (filter("inclass", QueryBuilders.termQuery("class", true))).subAggregation(
+                    significantTerms("mySignificantTerms").field("text")
+                        .minDocCount(2)
+                        .size(2)
+                        .shardSize(2)
+                        .executionHint(randomExecutionHint())
+                )
+            ),
+            response -> {
+                InternalFilter filteredBucket = response.getAggregations().get("inclass");
+                SignificantTerms sigterms = filteredBucket.getAggregations().get("mySignificantTerms");
+                assertThat(sigterms.getBuckets().size(), equalTo(0));
+            }
+        );
+
+        assertNoFailuresAndResponse(
+            prepareSearch(index).addAggregation(
+                (filter("inclass", QueryBuilders.termQuery("class", true))).subAggregation(
+                    significantTerms("mySignificantTerms").field("text")
+                        .minDocCount(2)
+                        .shardSize(2)
+                        .shardMinDocCount(2)
+                        .size(2)
+                        .executionHint(randomExecutionHint())
+                )
+            ),
+            response -> {
+                assertNoFailures(response);
+                InternalFilter filteredBucket = response.getAggregations().get("inclass");
+                SignificantTerms sigterms = filteredBucket.getAggregations().get("mySignificantTerms");
+                assertThat(sigterms.getBuckets().size(), equalTo(2));
+            }
+        );
     }
 
     private void addTermsDocs(String term, int numInClass, int numNotInClass, List<IndexRequestBuilder> builders) {
@@ -122,31 +129,36 @@ public class TermsShardMinDocCountIT extends ESIntegTestCase {
         indexRandom(true, false, indexBuilders);
 
         // first, check that indeed when not setting the shardMinDocCount parameter 0 terms are returned
-        SearchResponse response = prepareSearch(index).addAggregation(
-            terms("myTerms").field("text")
-                .minDocCount(2)
-                .size(2)
-                .shardSize(2)
-                .executionHint(randomExecutionHint())
-                .order(BucketOrder.key(true))
-        ).get();
-        assertNoFailures(response);
-        Terms sigterms = response.getAggregations().get("myTerms");
-        assertThat(sigterms.getBuckets().size(), equalTo(0));
-
-        response = prepareSearch(index).addAggregation(
-            terms("myTerms").field("text")
-                .minDocCount(2)
-                .shardMinDocCount(2)
-                .size(2)
-                .shardSize(2)
-                .executionHint(randomExecutionHint())
-                .order(BucketOrder.key(true))
-        ).get();
-        assertNoFailures(response);
-        sigterms = response.getAggregations().get("myTerms");
-        assertThat(sigterms.getBuckets().size(), equalTo(2));
-
+        assertNoFailuresAndResponse(
+            prepareSearch(index).addAggregation(
+                terms("myTerms").field("text")
+                    .minDocCount(2)
+                    .size(2)
+                    .shardSize(2)
+                    .executionHint(randomExecutionHint())
+                    .order(BucketOrder.key(true))
+            ),
+            response -> {
+                Terms sigterms = response.getAggregations().get("myTerms");
+                assertThat(sigterms.getBuckets().size(), equalTo(0));
+            }
+        );
+
+        assertNoFailuresAndResponse(
+            prepareSearch(index).addAggregation(
+                terms("myTerms").field("text")
+                    .minDocCount(2)
+                    .shardMinDocCount(2)
+                    .size(2)
+                    .shardSize(2)
+                    .executionHint(randomExecutionHint())
+                    .order(BucketOrder.key(true))
+            ),
+            response -> {
+                Terms sigterms = response.getAggregations().get("myTerms");
+                assertThat(sigterms.getBuckets().size(), equalTo(2));
+            }
+        );
     }
 
     private static void addTermsDocs(String term, int numDocs, List<IndexRequestBuilder> builders) {