Fix flaky test (#18235)

Signed-off-by: Prudhvi Godithi <[email protected]>

Author: prudhvigodithi

server/src/internalClusterTest/java/org/opensearch/search/simple/SimpleSearchIT.java

@@ -439,8 +439,14 @@ public void testSimpleTerminateAfterTrackTotalHitsUpToRandomSize0() throws Excep
         doTestSimpleTerminateAfterTrackTotalHitsUpTo(0);
     }
 
-    public void testSimpleTerminateAfterTrackTotalHitsUpToSize() throws Exception {
-        doTestSimpleTerminateAfterTrackTotalHitsUpTo(randomIntBetween(1, 29));
+    // Tests scenario where size <= track_total_hits_up_to (5)
+    public void testSimpleTerminateAfterTrackTotalHitsUpToSmallSize() throws Exception {
+        doTestSimpleTerminateAfterTrackTotalHitsUpTo(randomIntBetween(1, 5));
+    }
+
+    // Tests scenario where size > track_total_hits_up_to (5)
+    public void testSimpleTerminateAfterTrackTotalHitsUpToLargeSize() throws Exception {
+        doTestSimpleTerminateAfterTrackTotalHitsUpTo(randomIntBetween(6, 29));
     }
 
     public void testSimpleIndexSortEarlyTerminate() throws Exception {

server/src/main/java/org/opensearch/search/approximate/ApproximatePointRangeQuery.java

@@ -449,7 +449,11 @@ public boolean canApproximate(SearchContext context) {
         if (context.from() + context.size() == 0) {
             this.setSize(SearchContext.DEFAULT_TRACK_TOTAL_HITS_UP_TO);
         } else {
-            this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo()));
+            // We add +1 to ensure we collect at least one more document than required. This guarantees correct relation value:
+            // - If we find exactly trackTotalHitsUpTo docs: relation = EQUAL_TO
+            // - If we find > trackTotalHitsUpTo docs: relation = GREATER_THAN_OR_EQUAL_TO
+            // With +1, we will consistently get GREATER_THAN_OR_EQUAL_TO relation.
+            this.setSize(Math.max(context.from() + context.size(), context.trackTotalHitsUpTo()) + 1);
         }
         if (context.request() != null && context.request().source() != null) {
             FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(context.request().source());