Update Test Framework To Handle Query Rewrites That Rely on Non-Null Searchers (#129160)

test/framework/src/main/java/org/elasticsearch/test/AbstractQueryTestCase.java

@@ -9,10 +9,14 @@
 
 package org.elasticsearch.test;
 
+import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.index.RandomIndexWriter;
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.action.support.PlainActionFuture;
@@ -45,6 +49,7 @@ import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xcontent.json.JsonStringEncoder;
 import org.elasticsearch.xcontent.json.JsonXContent;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.time.Instant;
 import java.time.ZoneOffset;
@@ -453,13 +458,18 @@ public abstract class AbstractQueryTestCase<QB extends AbstractQueryBuilder<QB>>
         return true;
     }
 
+    protected IndexReaderManager getIndexReaderManager() {
+        return NullIndexReaderManager.INSTANCE;
+    }
+
     /**
      * Test creates the {@link Query} from the {@link QueryBuilder} under test and delegates the
      * assertions being made on the result to the implementing subclass.
      */
     public void testToQuery() throws IOException {
         for (int runs = 0; runs < NUMBER_OF_TESTQUERIES; runs++) {
-            SearchExecutionContext context = createSearchExecutionContext();
+            try (IndexReaderManager irm = getIndexReaderManager()) {
+                SearchExecutionContext context = createSearchExecutionContext(irm.getIndexSearcher());
                 assert context.isCacheable();
                 context.setAllowUnmappedFields(true);
                 QB firstQuery = createTestQueryBuilder();
@@ -501,9 +511,8 @@ public abstract class AbstractQueryTestCase<QB extends AbstractQueryBuilder<QB>>
                     );
                 }
                 context = new SearchExecutionContext(context);
-            Query secondLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context)).toQuery(
+                Query secondLuceneQuery = rewriteQuery(secondQuery, createQueryRewriteContext(), new SearchExecutionContext(context))
-                context
+                    .toQuery(context);
-            );
                 assertNotNull("toQuery should not return null", secondLuceneQuery);
                 assertLuceneQuery(secondQuery, secondLuceneQuery, context);
 
@@ -532,6 +541,7 @@ public abstract class AbstractQueryTestCase<QB extends AbstractQueryBuilder<QB>>
                 }
             }
         }
+    }
 
     /**
      * Simulate rewriting the query builder exclusively on the data node.
@@ -938,9 +948,75 @@ public abstract class AbstractQueryTestCase<QB extends AbstractQueryBuilder<QB>>
      */
     public void testCacheability() throws IOException {
         QB queryBuilder = createTestQueryBuilder();
-        SearchExecutionContext context = createSearchExecutionContext();
+        try (IndexReaderManager irm = getIndexReaderManager()) {
+            SearchExecutionContext context = createSearchExecutionContext(irm.getIndexSearcher());
             QueryBuilder rewriteQuery = rewriteQuery(queryBuilder, createQueryRewriteContext(), new SearchExecutionContext(context));
             assertNotNull(rewriteQuery.toQuery(context));
             assertTrue("query should be cacheable: " + queryBuilder.toString(), context.isCacheable());
         }
     }
+
+    public static class IndexReaderManager implements Closeable {
+        private final Directory directory;
+        private RandomIndexWriter indexWriter;
+        private IndexReader indexReader;
+        private IndexSearcher indexSearcher;
+
+        public IndexReaderManager() {
+            this.directory = newDirectory();
+        }
+
+        private IndexReaderManager(Directory directory) {
+            this.directory = directory;
+        }
+
+        public IndexReader getIndexReader() throws IOException {
+            if (indexReader == null) {
+                indexWriter = new RandomIndexWriter(random(), directory);
+                initIndexWriter(indexWriter);
+                indexReader = indexWriter.getReader();
+            }
+            return indexReader;
+        }
+
+        public IndexSearcher getIndexSearcher() throws IOException {
+            if (indexSearcher == null) {
+                indexSearcher = newSearcher(getIndexReader());
+            }
+            return indexSearcher;
+        }
+
+        @Override
+        public void close() throws IOException {
+            if (indexReader != null) {
+                indexReader.close();
+            }
+            if (indexWriter != null) {
+                indexWriter.close();
+            }
+            if (directory != null) {
+                directory.close();
+            }
+        }
+
+        protected void initIndexWriter(RandomIndexWriter indexWriter) {}
+    }
+
+    public static class NullIndexReaderManager extends IndexReaderManager {
+        public static final NullIndexReaderManager INSTANCE = new NullIndexReaderManager();
+
+        public NullIndexReaderManager() {
+            super(null);
+        }
+
+        @Override
+        public IndexReader getIndexReader() {
+            return null;
+        }
+
+        @Override
+        public IndexSearcher getIndexSearcher() {
+            return null;
+        }
+    }
+}