From a843bf43159c2188adc1043a42945bf784797254 Mon Sep 17 00:00:00 2001
From: Harsha Vamsi Kalluri <harshavamsi096@gmail.com>
Date: Tue, 19 Dec 2023 13:57:11 -0800
Subject: [PATCH] Fix tests for term and terms changes

Signed-off-by: Harsha Vamsi Kalluri <harshavamsi096@gmail.com>
---
 .../opensearch/index/mapper/BooleanFieldMapperTests.java   | 7 ++++++-
 .../org/opensearch/index/query/TermQueryBuilderTests.java  | 2 ++
 .../org/opensearch/index/query/TermsQueryBuilderTests.java | 2 ++
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/server/src/test/java/org/opensearch/index/mapper/BooleanFieldMapperTests.java b/server/src/test/java/org/opensearch/index/mapper/BooleanFieldMapperTests.java
index 8dec03a353d16..6981b1dd544be 100644
--- a/server/src/test/java/org/opensearch/index/mapper/BooleanFieldMapperTests.java
+++ b/server/src/test/java/org/opensearch/index/mapper/BooleanFieldMapperTests.java
@@ -32,12 +32,16 @@
 
 package org.opensearch.index.mapper;
 
+import org.apache.lucene.document.SortedNumericDocValuesField;
+import org.apache.lucene.document.SortedSetDocValuesField;
 import org.apache.lucene.index.DocValuesType;
 import org.apache.lucene.index.IndexableField;
 import org.apache.lucene.index.LeafReader;
 import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.index.SortedSetDocValues;
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.BoostQuery;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.util.BytesRef;
 import org.opensearch.common.xcontent.XContentFactory;
@@ -46,6 +50,7 @@
 import org.opensearch.index.mapper.ParseContext.Document;
 
 import java.io.IOException;
+import java.util.SortedSet;
 
 public class BooleanFieldMapperTests extends MapperTestCase {
 
@@ -206,7 +211,7 @@ public void testBoosts() throws Exception {
         }));
 
         MappedFieldType ft = mapperService.fieldType("field");
-        assertEquals(new BoostQuery(new TermQuery(new Term("field", "T")), 2.0f), ft.termQuery("true", null));
+        assertEquals(new IndexOrDocValuesQuery(new BoostQuery(new TermQuery(new Term("field", "T")), 2.0f), SortedNumericDocValuesField.newSlowExactQuery("field", 1)), ft.termQuery("true", null));
         assertParseMaximalWarnings();
     }
 }
diff --git a/server/src/test/java/org/opensearch/index/query/TermQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/TermQueryBuilderTests.java
index 3ac9bce840a20..c5bdf9b586df1 100644
--- a/server/src/test/java/org/opensearch/index/query/TermQueryBuilderTests.java
+++ b/server/src/test/java/org/opensearch/index/query/TermQueryBuilderTests.java
@@ -36,6 +36,7 @@
 
 import org.apache.lucene.index.Term;
 import org.apache.lucene.search.AutomatonQuery;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.PointRangeQuery;
 import org.apache.lucene.search.Query;
@@ -112,6 +113,7 @@ protected void doAssertLuceneQuery(TermQueryBuilder queryBuilder, Query query, Q
             either(instanceOf(TermQuery.class)).or(instanceOf(PointRangeQuery.class))
                 .or(instanceOf(MatchNoDocsQuery.class))
                 .or(instanceOf(AutomatonQuery.class))
+                .or(instanceOf(IndexOrDocValuesQuery.class))
         );
         MappedFieldType mapper = context.fieldMapper(queryBuilder.fieldName());
         if (query instanceof TermQuery) {
diff --git a/server/src/test/java/org/opensearch/index/query/TermsQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/TermsQueryBuilderTests.java
index b587bec2d5343..32bf290627b63 100644
--- a/server/src/test/java/org/opensearch/index/query/TermsQueryBuilderTests.java
+++ b/server/src/test/java/org/opensearch/index/query/TermsQueryBuilderTests.java
@@ -34,6 +34,7 @@
 
 import org.apache.lucene.search.BooleanQuery;
 import org.apache.lucene.search.ConstantScoreQuery;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.PointInSetQuery;
 import org.apache.lucene.search.Query;
@@ -135,6 +136,7 @@ protected void doAssertLuceneQuery(TermsQueryBuilder queryBuilder, Query query,
                 either(instanceOf(TermInSetQuery.class)).or(instanceOf(PointInSetQuery.class))
                     .or(instanceOf(ConstantScoreQuery.class))
                     .or(instanceOf(MatchNoDocsQuery.class))
+                    .or(instanceOf(IndexOrDocValuesQuery.class))
             );
             if (query instanceof ConstantScoreQuery) {
                 assertThat(((ConstantScoreQuery) query).getQuery(), instanceOf(BooleanQuery.class));