From db9b469d58cc2b7beff63e51d3005de4fbb9f228 Mon Sep 17 00:00:00 2001 From: Panagiotis Bailis Date: Tue, 10 Sep 2024 16:32:18 +0300 Subject: [PATCH] iter --- .../RankDocsRetrieverBuilderTests.java | 24 ++++++++++++++----- 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilderTests.java b/server/src/test/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilderTests.java index 1473540fea5ef..d1de451253fb2 100644 --- a/server/src/test/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilderTests.java +++ b/server/src/test/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilderTests.java @@ -18,9 +18,9 @@ import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder; import org.elasticsearch.search.builder.SearchSourceBuilder; import org.elasticsearch.search.rank.RankDoc; +import org.elasticsearch.search.rank.RankDocsRankBuilder; import org.elasticsearch.search.retriever.rankdoc.RankDocsQueryBuilder; import org.elasticsearch.search.retriever.rankdoc.RankDocsSortBuilder; -import org.elasticsearch.search.sort.ScoreSortBuilder; import org.elasticsearch.test.ESTestCase; import java.io.IOException; @@ -100,14 +100,26 @@ public void testExtractToSearchSourceBuilder() { if (randomBoolean()) { source.aggregation(new TermsAggregationBuilder("name").field("field")); } + if (randomBoolean()) { + source.explain(true); + } + if (randomBoolean()) { + source.profile(true); + } retriever.extractToSearchSourceBuilder(source, randomBoolean()); - assertThat(source.sorts().size(), equalTo(2)); + assertThat(source.sorts().size(), equalTo(1)); assertThat(source.sorts().get(0), instanceOf(RankDocsSortBuilder.class)); - assertThat(source.sorts().get(1), instanceOf(ScoreSortBuilder.class)); + assertThat(source.rankBuilder(), instanceOf(RankDocsRankBuilder.class)); assertThat(source.query(), instanceOf(BoolQueryBuilder.class)); BoolQueryBuilder bq = (BoolQueryBuilder) source.query(); - if (source.aggregations() != null) { - assertThat(bq.must().size(), equalTo(0)); + if (source.aggregations() != null || (source.explain() != null && source.explain()) || source.profile()) { + if (source.profile() || (source.explain() != null && source.explain())) { + assertThat(bq.must().size(), equalTo(retriever.sources.size())); + assertThat(bq.filter().size(), equalTo(retriever.preFilterQueryBuilders.size())); + } else { + assertThat(bq.must().size(), equalTo(0)); + assertThat(bq.filter().size(), equalTo(retriever.preFilterQueryBuilders.size() + retriever.sources.size())); + } assertThat(bq.should().size(), greaterThanOrEqualTo(1)); assertThat(bq.should().get(0), instanceOf(RankDocsQueryBuilder.class)); assertNull(source.postFilter()); @@ -115,8 +127,8 @@ public void testExtractToSearchSourceBuilder() { assertThat(bq.must().size(), equalTo(1)); assertThat(bq.must().get(0), instanceOf(RankDocsQueryBuilder.class)); assertNull(source.postFilter()); + assertThat(bq.filter().size(), equalTo(retriever.preFilterQueryBuilders.size())); } - assertThat(bq.filter().size(), equalTo(retriever.preFilterQueryBuilders.size())); } public void testTopDocsQuery() {