Fix SearchExecutionContextTests.testFielddataLookupSomtimesLoop

server/src/test/java/org/elasticsearch/index/query/SearchExecutionContextTests.java

@@ -10,7 +10,6 @@
 
 import org.apache.lucene.document.Field;
 import org.apache.lucene.document.KeywordField;
-import org.apache.lucene.document.StringField;
 import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.index.Term;
@@ -223,18 +222,27 @@ public void testFielddataLookupTerminatesInLoop() {
     }
 
     public void testFielddataLookupSometimesLoop() throws IOException {
-        SearchExecutionContext searchExecutionContext = createSearchExecutionContext(
-            // simulate a runtime field cycle in the second doc: 1: doc['2'] 2: doc['3'] 3: doc['4'] 4: doc['4']
+        // create this field so we can use it to make sure we're escaping the loop on only the "first" document
+        var concreteField = new KeywordFieldMapper.KeywordFieldType("indexed_field", true, true, Collections.emptyMap());
+
+        // simulate a runtime field cycle in the second doc: 1: doc['2'] 2: doc['3'] 3: doc['4'] 4: doc['4']
+        var runtimeFields = List.of(
             runtimeField("1", leafLookup -> leafLookup.doc().get("2").get(0).toString()),
             runtimeField("2", leafLookup -> leafLookup.doc().get("3").get(0).toString()),
             runtimeField("3", leafLookup -> leafLookup.doc().get("4").get(0).toString()),
-            runtimeField("4", (leafLookup, docId) -> {
-                if (docId == 0) {
+            runtimeField("4", leafLookup -> {
+                if (leafLookup.doc().get("indexed_field").getFirst().equals("first")) {
                     return "escape!";
                 }
-                return leafLookup.doc().get("4").get(0).toString();
+                return leafLookup.doc().get("4").getFirst().toString();
             })
         );
+        SearchExecutionContext searchExecutionContext = createSearchExecutionContext(
+            "uuid",
+            null,
+            createMappingLookup(List.of(concreteField), runtimeFields),
+            Collections.emptyMap()
+        );
         List<String> values = collect("1", searchExecutionContext, new TermQuery(new Term("indexed_field", "first")));
         assertEquals(List.of("escape!"), values);
         IllegalArgumentException iae = expectThrows(IllegalArgumentException.class, () -> collect("1", searchExecutionContext));
@@ -776,8 +784,8 @@ private static List<String> collect(String field, SearchExecutionContext searchE
     private static List<String> collect(String field, SearchExecutionContext searchExecutionContext, Query query) throws IOException {
         List<String> result = new ArrayList<>();
         try (Directory directory = newDirectory(); RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) {
-            indexWriter.addDocument(List.of(new StringField("indexed_field", "first", Field.Store.NO)));
-            indexWriter.addDocument(List.of(new StringField("indexed_field", "second", Field.Store.NO)));
+            indexWriter.addDocument(List.of(new KeywordField("indexed_field", "first", Field.Store.YES)));
+            indexWriter.addDocument(List.of(new KeywordField("indexed_field", "second", Field.Store.YES)));
             try (DirectoryReader reader = indexWriter.getReader()) {
                 IndexSearcher searcher = newSearcher(reader);
                 MappedFieldType fieldType = searchExecutionContext.getFieldType(field);