Unit test with scores

Author: jan-elastic

server/src/test/java/org/elasticsearch/search/aggregations/bucket/sampler/random/RandomSamplerAggregatorTests.java

@@ -11,22 +11,29 @@
 
 import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.document.SortedNumericDocValuesField;
+import org.apache.lucene.index.Term;
+import org.apache.lucene.search.BooleanClause;
+import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.index.mapper.KeywordFieldMapper;
 import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.aggregations.AggregationBuilders;
 import org.elasticsearch.search.aggregations.AggregatorTestCase;
 import org.elasticsearch.search.aggregations.bucket.filter.Filter;
 import org.elasticsearch.search.aggregations.metrics.Avg;
 import org.elasticsearch.search.aggregations.metrics.Max;
 import org.elasticsearch.search.aggregations.metrics.Min;
+import org.elasticsearch.search.aggregations.metrics.TopHits;
 import org.hamcrest.Description;
 import org.hamcrest.Matcher;
 import org.hamcrest.TypeSafeMatcher;
 
 import java.io.IOException;
+import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.DoubleStream;
@@ -37,6 +44,8 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
+import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.lessThan;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.hamcrest.Matchers.not;
 import static org.hamcrest.Matchers.notANumber;
@@ -76,6 +85,35 @@ public void testAggregationSampling() throws IOException {
         assertThat(avgAvg, closeTo(1.5, 0.5));
     }
 
+    public void testAggregationSampling_withScores() throws IOException {
+        long[] counts = new long[5];
+        AtomicInteger integer = new AtomicInteger();
+        do {
+            testCase(RandomSamplerAggregatorTests::writeTestDocs, (InternalRandomSampler result) -> {
+                counts[integer.get()] = result.getDocCount();
+                if (result.getDocCount() > 0) {
+                    TopHits agg = result.getAggregations().get("top");
+                    List<SearchHit> hits = Arrays.asList(agg.getHits().getHits());
+                    assertThat(Strings.toString(result), hits, hasSize(1));
+                    assertThat(Strings.toString(result), hits.get(0).getScore(), allOf(greaterThan(0.0f), lessThan(1.0f)));
+                }
+            },
+                new AggTestConfig(
+                    new RandomSamplerAggregationBuilder("my_agg").subAggregation(AggregationBuilders.topHits("top").size(1))
+                        .setProbability(0.25),
+                    longField(NUMERIC_FIELD_NAME)
+                ).withQuery(
+                    new BooleanQuery.Builder().add(
+                        new TermQuery(new Term(KEYWORD_FIELD_NAME, KEYWORD_FIELD_VALUE)),
+                        BooleanClause.Occur.SHOULD
+                    ).build()
+                )
+            );
+        } while (integer.incrementAndGet() < 5);
+        long avgCount = LongStream.of(counts).sum() / integer.get();
+        assertThat(avgCount, allOf(greaterThanOrEqualTo(20L), lessThanOrEqualTo(70L)));
+    }
+
     public void testAggregationSamplingNestedAggsScaled() throws IOException {
         // in case 0 docs get sampled, which can rarely happen
         // in case the test index has many segments.