Fix more tests

Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
harshavamsi · Jun 11, 2024 · 07f7d05 · 07f7d05
1 parent 6a28cc1
commit 07f7d05
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 5 deletions.
diff --git a/server/src/test/java/org/opensearch/index/mapper/RangeFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/RangeFieldTypeTests.java
@@ -57,6 +57,7 @@
 import org.opensearch.index.mapper.RangeFieldMapper.RangeFieldType;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.QueryShardException;
+import org.opensearch.search.approximate.ApproximateableQuery;
 import org.opensearch.test.IndexSettingsModule;
 import org.joda.time.DateTime;
 import org.junit.Before;
@@ -285,7 +286,10 @@ public void testDateRangeQueryUsingMappingFormatLegacy() {
         // compare lower and upper bounds with what we would get on a `date` field
         DateFieldType dateFieldType = new DateFieldType("field", DateFieldMapper.Resolution.MILLISECONDS, formatter);
         final Query queryOnDateField = dateFieldType.rangeQuery(from, to, true, true, relation, null, fieldType.dateMathParser(), context);
-        assertEquals("field:[1465975790000 TO 1466062190999]", ((IndexOrDocValuesQuery) queryOnDateField).getIndexQuery().toString());
+        assertEquals(
+            "field:[1465975790000 TO 1466062190999]",
+            ((ApproximateableQuery) ((IndexOrDocValuesQuery) queryOnDateField).getIndexQuery()).getOriginalQuery().toString()
+        );
     }
 
     public void testDateRangeQueryUsingMappingFormat() {

diff --git a/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java b/server/src/test/java/org/opensearch/index/query/QueryStringQueryBuilderTests.java
@@ -53,6 +53,7 @@
 import org.apache.lucene.search.MultiTermQuery;
 import org.apache.lucene.search.NormsFieldExistsQuery;
 import org.apache.lucene.search.PhraseQuery;
+import org.apache.lucene.search.PointRangeQuery;
 import org.apache.lucene.search.PrefixQuery;
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.RegexpQuery;
@@ -76,6 +77,8 @@
 import org.opensearch.index.mapper.FieldNamesFieldMapper;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.search.QueryStringQueryParser;
+import org.opensearch.search.approximate.ApproximatePointRangeQuery;
+import org.opensearch.search.approximate.ApproximateableQuery;
 import org.opensearch.test.AbstractQueryTestCase;
 import org.hamcrest.CoreMatchers;
 import org.hamcrest.Matchers;
@@ -98,6 +101,7 @@
 import static org.hamcrest.CoreMatchers.hasItems;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.instanceOf;
+import static org.apache.lucene.document.LongPoint.pack;
 
 public class QueryStringQueryBuilderTests extends AbstractQueryTestCase<QueryStringQueryBuilder> {
 
@@ -856,14 +860,42 @@ public void testToQueryDateWithTimeZone() throws Exception {
         assertThat(query, instanceOf(IndexOrDocValuesQuery.class));
         long lower = 0; // 1970-01-01T00:00:00.999 UTC
         long upper = 86399999;  // 1970-01-01T23:59:59.999 UTC
-        assertEquals(calculateExpectedDateQuery(lower, upper), query);
+        assertEquals(calculateExpectedDateQuery(lower, upper).toString(), query.toString());
         int msPerHour = 3600000;
-        assertEquals(calculateExpectedDateQuery(lower - msPerHour, upper - msPerHour), qsq.timeZone("+01:00").toQuery(context));
-        assertEquals(calculateExpectedDateQuery(lower + msPerHour, upper + msPerHour), qsq.timeZone("-01:00").toQuery(context));
+        assertEquals(
+            calculateExpectedDateQuery(lower - msPerHour, upper - msPerHour).toString(),
+            qsq.timeZone("+01:00").toQuery(context).toString()
+        );
+        assertEquals(
+            calculateExpectedDateQuery(lower + msPerHour, upper + msPerHour).toString(),
+            qsq.timeZone("-01:00").toQuery(context).toString()
+        );
     }
 
     private IndexOrDocValuesQuery calculateExpectedDateQuery(long lower, long upper) {
-        Query query = LongPoint.newRangeQuery(DATE_FIELD_NAME, lower, upper);
+        Query query = new ApproximateableQuery(
+            new PointRangeQuery(
+                DATE_FIELD_NAME,
+                pack(new long[] { lower }).bytes,
+                pack(new long[] { upper }).bytes,
+                new long[] { lower }.length
+            ) {
+                protected String toString(int dimension, byte[] value) {
+                    return Long.toString(LongPoint.decodeDimension(value, 0));
+                }
+            },
+            new ApproximatePointRangeQuery(
+                DATE_FIELD_NAME,
+                pack(new long[] { lower }).bytes,
+                pack(new long[] { upper }).bytes,
+                new long[] { lower }.length
+            ) {
+                @Override
+                protected String toString(int dimension, byte[] value) {
+                    return Long.toString(LongPoint.decodeDimension(value, 0));
+                }
+            }
+        );
         Query dv = SortedNumericDocValuesField.newSlowRangeQuery(DATE_FIELD_NAME, lower, upper);
         return new IndexOrDocValuesQuery(query, dv);
     }