Refactor tests

Author: carlosdelest

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryEvaluatorTests.java

@@ -25,6 +25,8 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.compute.OperatorTests;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.BytesRefVector;
 import org.elasticsearch.compute.data.DocBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.ElementType;
@@ -50,6 +52,7 @@
 import java.util.TreeSet;
 
 import static org.elasticsearch.compute.test.OperatorTestCase.randomPageSize;
+import static org.hamcrest.Matchers.equalTo;
 
 /**
  * Base class for testing Lucene query evaluators.
@@ -149,7 +152,26 @@ private void testTermsQuery(boolean shuffleDocs) throws IOException {
         assertTermsQuery(results, matching, expectedMatchCount);
     }
 
-    protected abstract void assertTermsQuery(List<Page> results, Set<String> matching, int expectedMatchCount);
+    protected void assertTermsQuery(List<Page> results, Set<String> matching, int expectedMatchCount) {
+        int matchCount = 0;
+        for (Page page : results) {
+            int initialBlockIndex = termsBlockIndex(page);
+            BytesRefVector terms = page.<BytesRefBlock>getBlock(initialBlockIndex).asVector();
+            @SuppressWarnings("unchecked")
+            T resultVector = (T) page.getBlock(resultsBlockIndex(page)).asVector();
+            for (int i = 0; i < page.getPositionCount(); i++) {
+                BytesRef termAtPosition = terms.getBytesRef(i, new BytesRef());
+                boolean isMatch = matching.contains(termAtPosition.utf8ToString());
+                assertResultMatch(resultVector, i, isMatch);
+                if (isMatch) {
+                    matchCount++;
+                }
+            }
+        }
+        assertThat(matchCount, equalTo(expectedMatchCount));
+    }
+
+    protected abstract void assertResultMatch(T resultVector, int position, boolean isMatch);
 
     private List<Page> runQuery(Set<String> values, Query query, boolean shuffleDocs) throws IOException {
         DriverContext driverContext = driverContext();

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryExpressionEvaluatorTests.java

@@ -78,20 +78,7 @@ protected int resultsBlockIndex(Page page) {
     }
 
     @Override
-    protected void assertTermsQuery(List<Page> results, Set<String> matching, int expectedMatchCount) {
-        int matchCount = 0;
-        for (Page page : results) {
-            int initialBlockIndex = termsBlockIndex(page);
-            BytesRefVector terms = page.<BytesRefBlock>getBlock(initialBlockIndex).asVector();
-            BooleanVector matches = page.<BooleanBlock>getBlock(initialBlockIndex + 1).asVector();
-            for (int i = 0; i < page.getPositionCount(); i++) {
-                BytesRef termAtPosition = terms.getBytesRef(i, new BytesRef());
-                assertThat(matches.getBoolean(i), equalTo(matching.contains(termAtPosition.utf8ToString())));
-                if (matches.getBoolean(i)) {
-                    matchCount++;
-                }
-            }
-        }
-        assertThat(matchCount, equalTo(expectedMatchCount));
+    protected void assertResultMatch(BooleanVector resultVector, int position, boolean isMatch) {
+        assertThat(resultVector.getBoolean(position), equalTo(isMatch));
     }
 }

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/lucene/LuceneQueryScoreEvaluatorTests.java

@@ -8,18 +8,13 @@
 package org.elasticsearch.compute.lucene;
 
 import org.apache.lucene.search.Scorable;
-import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.compute.data.BlockFactory;
-import org.elasticsearch.compute.data.BytesRefBlock;
-import org.elasticsearch.compute.data.BytesRefVector;
 import org.elasticsearch.compute.data.DoubleVector;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.compute.operator.ScoreOperator;
 
 import java.io.IOException;
-import java.util.List;
-import java.util.Set;
 
 import static org.elasticsearch.compute.lucene.LuceneQueryScoreEvaluator.NO_MATCH_SCORE;
 import static org.hamcrest.Matchers.equalTo;
@@ -82,23 +77,12 @@ protected int resultsBlockIndex(Page page) {
     }
 
     @Override
-    protected void assertTermsQuery(List<Page> results, Set<String> matching, int expectedMatchCount) {
-        int matchCount = 0;
-        for (Page page : results) {
-            int initialBlockIndex = termsBlockIndex(page);
-            BytesRefVector terms = page.<BytesRefBlock>getBlock(initialBlockIndex).asVector();
-            DoubleVector matches = (DoubleVector) page.getBlock(resultsBlockIndex(page)).asVector();
-            for (int i = 0; i < page.getPositionCount(); i++) {
-                BytesRef termAtPosition = terms.getBytesRef(i, new BytesRef());
-                if (matching.contains(termAtPosition.utf8ToString())) {
-                    assertThat(matches.getDouble(i), greaterThan((double) TEST_SCORE));
-                    matchCount++;
-                } else {
-                    // Default score, as Lucene docs gets retrieved with a implicit score
-                    assertThat(matches.getDouble(i), equalTo(DEFAULT_SCORE));
-                }
-            }
+    protected void assertResultMatch(DoubleVector resultVector, int position, boolean isMatch) {
+        if (isMatch) {
+            assertThat(resultVector.getDouble(position), greaterThan(DEFAULT_SCORE));
+        } else {
+            // All docs have a default score coming from Lucene
+            assertThat(resultVector.getDouble(position), equalTo(DEFAULT_SCORE));
         }
-        assertThat(matchCount, equalTo(expectedMatchCount));
     }
 }