From ae85095159875a805dfb5d4e8f6ea622cdc6846e Mon Sep 17 00:00:00 2001
From: Harsha Vamsi Kalluri <harshavamsi096@gmail.com>
Date: Thu, 8 Aug 2024 10:17:12 -0700
Subject: [PATCH] Add overrides to more places

Signed-off-by: Harsha Vamsi Kalluri <harshavamsi096@gmail.com>
---
 .../mapper/ScaledFloatFieldTypeTests.java     |  1 +
 .../index/mapper/NumberFieldTypeTests.java    | 51 +++++++++++++++----
 2 files changed, 43 insertions(+), 9 deletions(-)

diff --git a/modules/mapper-extras/src/test/java/org/opensearch/index/mapper/ScaledFloatFieldTypeTests.java b/modules/mapper-extras/src/test/java/org/opensearch/index/mapper/ScaledFloatFieldTypeTests.java
index a653edbd05992..b0f4ea73cf9e5 100644
--- a/modules/mapper-extras/src/test/java/org/opensearch/index/mapper/ScaledFloatFieldTypeTests.java
+++ b/modules/mapper-extras/src/test/java/org/opensearch/index/mapper/ScaledFloatFieldTypeTests.java
@@ -122,6 +122,7 @@ public void testRangeQuery() throws IOException {
                 u,
                 includeLower,
                 includeUpper,
+                RewriteOverride.DEFAULT,
                 false,
                 true,
                 MOCK_QSC
diff --git a/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java b/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java
index 96487db6dd512..b6769739945db 100644
--- a/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java
+++ b/server/src/test/java/org/opensearch/index/mapper/NumberFieldTypeTests.java
@@ -550,8 +550,28 @@ public void testHalfFloatRange() throws IOException {
             float u = (randomFloat() * 2 - 1) * 65504;
             boolean includeLower = randomBoolean();
             boolean includeUpper = randomBoolean();
-            Query floatQ = NumberType.FLOAT.rangeQuery("float", l, u, includeLower, includeUpper, false, true, MOCK_QSC);
-            Query halfFloatQ = NumberType.HALF_FLOAT.rangeQuery("half_float", l, u, includeLower, includeUpper, false, true, MOCK_QSC);
+            Query floatQ = NumberType.FLOAT.rangeQuery(
+                "float",
+                l,
+                u,
+                includeLower,
+                includeUpper,
+                RewriteOverride.DEFAULT,
+                false,
+                true,
+                MOCK_QSC
+            );
+            Query halfFloatQ = NumberType.HALF_FLOAT.rangeQuery(
+                "half_float",
+                l,
+                u,
+                includeLower,
+                includeUpper,
+                RewriteOverride.DEFAULT,
+                false,
+                true,
+                MOCK_QSC
+            );
             assertEquals(searcher.count(floatQ), searcher.count(halfFloatQ));
         }
         IOUtils.close(reader, dir);
@@ -587,11 +607,22 @@ public void testUnsignedLongRange() throws IOException {
                 u,
                 includeLower,
                 includeUpper,
+                RewriteOverride.DEFAULT,
+                false,
+                true,
+                MOCK_QSC
+            );
+            Query doubleQ = NumberType.DOUBLE.rangeQuery(
+                "double",
+                l,
+                u,
+                includeLower,
+                includeUpper,
+                RewriteOverride.DEFAULT,
                 false,
                 true,
                 MOCK_QSC
             );
-            Query doubleQ = NumberType.DOUBLE.rangeQuery("double", l, u, includeLower, includeUpper, false, true, MOCK_QSC);
             assertEquals(searcher.count(doubleQ), searcher.count(unsignedLongQ));
         }
         IOUtils.close(reader, dir);
@@ -599,16 +630,16 @@ public void testUnsignedLongRange() throws IOException {
 
     public void testNegativeZero() {
         assertEquals(
-            NumberType.DOUBLE.rangeQuery("field", null, -0d, true, true, false, true, MOCK_QSC),
-            NumberType.DOUBLE.rangeQuery("field", null, +0d, true, false, false, true, MOCK_QSC)
+            NumberType.DOUBLE.rangeQuery("field", null, -0d, true, true, RewriteOverride.DEFAULT, false, true, MOCK_QSC),
+            NumberType.DOUBLE.rangeQuery("field", null, +0d, true, false, RewriteOverride.DEFAULT, false, true, MOCK_QSC)
         );
         assertEquals(
-            NumberType.FLOAT.rangeQuery("field", null, -0f, true, true, false, true, MOCK_QSC),
-            NumberType.FLOAT.rangeQuery("field", null, +0f, true, false, false, true, MOCK_QSC)
+            NumberType.FLOAT.rangeQuery("field", null, -0f, true, true, RewriteOverride.DEFAULT, false, true, MOCK_QSC),
+            NumberType.FLOAT.rangeQuery("field", null, +0f, true, false, RewriteOverride.DEFAULT, false, true, MOCK_QSC)
         );
         assertEquals(
-            NumberType.HALF_FLOAT.rangeQuery("field", null, -0f, true, true, false, true, MOCK_QSC),
-            NumberType.HALF_FLOAT.rangeQuery("field", null, +0f, true, false, false, true, MOCK_QSC)
+            NumberType.HALF_FLOAT.rangeQuery("field", null, -0f, true, true, RewriteOverride.DEFAULT, false, true, MOCK_QSC),
+            NumberType.HALF_FLOAT.rangeQuery("field", null, +0f, true, false, RewriteOverride.DEFAULT, false, true, MOCK_QSC)
         );
 
         assertFalse(NumberType.DOUBLE.termQuery("field", -0d, true, true).equals(NumberType.DOUBLE.termQuery("field", +0d, true, true)));
@@ -670,6 +701,7 @@ public void doTestDocValueRangeQueries(NumberType type, Supplier<Number> valueSu
                 random().nextBoolean() ? null : valueSupplier.get(),
                 randomBoolean(),
                 randomBoolean(),
+                RewriteOverride.DEFAULT,
                 true,
                 true,
                 MOCK_QSC
@@ -753,6 +785,7 @@ public void doTestIndexSortRangeQueries(NumberType type, Supplier<Number> valueS
                 random().nextBoolean() ? null : valueSupplier.get(),
                 randomBoolean(),
                 randomBoolean(),
+                RewriteOverride.DEFAULT,
                 true,
                 true,
                 context