Changing to modified term query in test

cbuescher · Jul 5, 2023 · 40e8f00 · 40e8f00
1 parent ffa7ea7
commit 40e8f00
Showing 1 changed file with 22 additions and 22 deletions.
diff --git a/server/src/test/java/org/elasticsearch/search/query/PartialHitCountCollectorTests.java b/server/src/test/java/org/elasticsearch/search/query/PartialHitCountCollectorTests.java
@@ -16,6 +16,7 @@
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BooleanClause;
 import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.FilterWeight;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.Query;
@@ -27,7 +28,6 @@
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
-import java.util.List;
 
 public class PartialHitCountCollectorTests extends ESTestCase {
 
@@ -105,33 +105,33 @@ public void testCollectedHitCount() throws Exception {
     }
 
     public void testCollectedHitCountEarlyTerminated() throws Exception {
-        Query query = new BooleanQuery.Builder().add(new TermQuery(new Term("string", "a1")), BooleanClause.Occur.SHOULD)
-            .add(new TermQuery(new Term("string", "b3")), BooleanClause.Occur.SHOULD)
-            .add(new TermQuery(new Term("string", "b6")), BooleanClause.Occur.SHOULD)
+        Query query = new BooleanQuery.Builder().add(termQuery(new Term("string", "a1")), BooleanClause.Occur.SHOULD)
+            .add(termQuery(new Term("string", "b3")), BooleanClause.Occur.SHOULD)
+            .add(termQuery(new Term("string", "b6")), BooleanClause.Occur.SHOULD)
             .build();
         // there's three docs matching the query: any totalHitsThreshold lower than 3 will trigger early termination
         int totalHitsThreshold = randomInt(2);
         PartialHitCountCollector partialHitCountCollector = new PartialHitCountCollector(totalHitsThreshold);
         searcher.search(query, partialHitCountCollector);
 
-        try {
-            assertEquals(totalHitsThreshold, partialHitCountCollector.getTotalHits());
-            assertTrue(partialHitCountCollector.hasEarlyTerminated());
-        } catch (AssertionError e) {
-            // if we get very unlucky with the document distribution, we can get more than one segment with only one document
-            // in this case early termination might not kick in because we never collect
-            // check if this was the case here and disregard error in that case
-            Weight weight = searcher.createWeight(query, ScoreMode.TOP_DOCS, 1.0f);
-            List<LeafReaderContext> leafContexts = searcher.getLeafContexts();
-            int totalHitsWithoutCollecting = 0;
-            for (LeafReaderContext leafContext : leafContexts) {
-                int count = weight.count(leafContext);
-                if (count != -1) {
-                    totalHitsWithoutCollecting += count;
-                }
+        assertEquals(totalHitsThreshold, partialHitCountCollector.getTotalHits());
+        assertTrue(partialHitCountCollector.hasEarlyTerminated());
+    }
+
+    /**
+     * Create a {@link TermQuery} which cannot terminate collection early
+     */
+    private TermQuery termQuery(Term term) {
+        return new TermQuery(term) {
+            public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+
+                Weight w = super.createWeight(searcher, scoreMode, boost);
+                return new FilterWeight(w) {
+                    public int count(LeafReaderContext context) throws IOException {
+                        return -1;
+                    }
+                };
             }
-            assertTrue(totalHitsWithoutCollecting > totalHitsThreshold);
-            assertFalse(partialHitCountCollector.hasEarlyTerminated());
-        }
+        };
     }
 }