| 
10 | 10 | package org.elasticsearch.search.aggregations.bucket;  | 
11 | 11 | 
 
  | 
12 | 12 | import org.elasticsearch.action.index.IndexRequestBuilder;  | 
 | 13 | +import org.elasticsearch.action.search.SearchRequestBuilder;  | 
13 | 14 | import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;  | 
14 | 15 | import org.elasticsearch.search.aggregations.bucket.sampler.random.InternalRandomSampler;  | 
15 | 16 | import org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplerAggregationBuilder;  | 
 | 
20 | 21 | import java.util.HashMap;  | 
21 | 22 | import java.util.List;  | 
22 | 23 | import java.util.Map;  | 
 | 24 | +import java.util.stream.IntStream;  | 
23 | 25 | 
 
  | 
24 | 26 | import static org.elasticsearch.search.aggregations.AggregationBuilders.avg;  | 
25 | 27 | import static org.elasticsearch.search.aggregations.AggregationBuilders.histogram;  | 
26 | 28 | import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertNoFailures;  | 
27 | 29 | import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponse;  | 
 | 30 | +import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertResponses;  | 
28 | 31 | import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;  | 
29 | 32 | import static org.hamcrest.Matchers.lessThan;  | 
30 | 33 | 
 
  | 
@@ -112,27 +115,28 @@ public void testRandomSamplerConsistentSeed() {  | 
112 | 115 |             }  | 
113 | 116 |         );  | 
114 | 117 | 
 
  | 
115 |  | -        for (int i = 0; i < NUM_SAMPLE_RUNS; i++) {  | 
116 |  | -            assertResponse(  | 
117 |  | -                prepareSearch("idx").setPreference("shard:0")  | 
118 |  | -                    .addAggregation(  | 
119 |  | -                        new RandomSamplerAggregationBuilder("sampler").setProbability(PROBABILITY)  | 
120 |  | -                            .setSeed(0)  | 
121 |  | -                            .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))  | 
122 |  | -                            .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))  | 
123 |  | -                            .setShardSeed(42)  | 
124 |  | -                    ),  | 
125 |  | -                response -> {  | 
126 |  | -                    InternalRandomSampler sampler = response.getAggregations().get("sampler");  | 
127 |  | -                    double monotonicValue = ((Avg) sampler.getAggregations().get("mean_monotonic")).getValue();  | 
128 |  | -                    double numericValue = ((Avg) sampler.getAggregations().get("mean_numeric")).getValue();  | 
129 |  | -                    long docCount = sampler.getDocCount();  | 
130 |  | -                    assertEquals(monotonicValue, sampleMonotonicValue[0], tolerance);  | 
131 |  | -                    assertEquals(numericValue, sampleNumericValue[0], tolerance);  | 
132 |  | -                    assertEquals(docCount, sampledDocCount[0]);  | 
133 |  | -                }  | 
134 |  | -            );  | 
135 |  | -        }  | 
 | 118 | +        assertResponses(response -> {  | 
 | 119 | +            InternalRandomSampler sampler = response.getAggregations().get("sampler");  | 
 | 120 | +            double monotonicValue = ((Avg) sampler.getAggregations().get("mean_monotonic")).getValue();  | 
 | 121 | +            double numericValue = ((Avg) sampler.getAggregations().get("mean_numeric")).getValue();  | 
 | 122 | +            long docCount = sampler.getDocCount();  | 
 | 123 | +            assertEquals(monotonicValue, sampleMonotonicValue[0], tolerance);  | 
 | 124 | +            assertEquals(numericValue, sampleNumericValue[0], tolerance);  | 
 | 125 | +            assertEquals(docCount, sampledDocCount[0]);  | 
 | 126 | +        },  | 
 | 127 | +            IntStream.rangeClosed(0, NUM_SAMPLE_RUNS - 1)  | 
 | 128 | +                .mapToObj(  | 
 | 129 | +                    num -> prepareSearch("idx").setPreference("shard:0")  | 
 | 130 | +                        .addAggregation(  | 
 | 131 | +                            new RandomSamplerAggregationBuilder("sampler").setProbability(PROBABILITY)  | 
 | 132 | +                                .setSeed(0)  | 
 | 133 | +                                .subAggregation(avg("mean_monotonic").field(MONOTONIC_VALUE))  | 
 | 134 | +                                .subAggregation(avg("mean_numeric").field(NUMERIC_VALUE))  | 
 | 135 | +                                .setShardSeed(42)  | 
 | 136 | +                        )  | 
 | 137 | +                )  | 
 | 138 | +                .toArray(SearchRequestBuilder[]::new)  | 
 | 139 | +        );  | 
136 | 140 |     }  | 
137 | 141 | 
 
  | 
138 | 142 |     public void testRandomSampler() {  | 
 | 
0 commit comments