harden tests

Author: dnhatn

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/ReplaceRoundToWithQueryAndTagsTests.java

@@ -69,7 +69,6 @@
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateNanosToLong;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.dateTimeToLong;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.is;
 
 //@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug")
@@ -540,61 +539,71 @@ public void testRoundToTransformToQueryAndTagsWithCustomizedUpperLimit() {
         }
     }
 
-    static String points(int numPoints) {
+    static String pointArray(int numPoints) {
         return IntStream.range(0, numPoints).mapToObj(Integer::toString).collect(Collectors.joining(","));
     }
 
+    static int queryAndTags(PhysicalPlan plan) {
+        EsQueryExec esQuery = (EsQueryExec) plan.collectFirstChildren(EsQueryExec.class::isInstance).getFirst();
+        return esQuery.queryBuilderAndTags().size();
+    }
+
     public void testAdjustThresholdForQueries() {
         {
+            int points = between(2, 127);
             String q = String.format(Locale.ROOT, """
                 from test
                 | stats count(*) by x = round_to(integer, %s)
-                """, points(between(1, 128)));
+                """, pointArray(points));
             PhysicalPlan plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
-            EsQueryExec esQuery = (EsQueryExec) plan.collectFirstChildren(EsQueryExec.class::isInstance).getFirst();
-            assertThat(esQuery.queryBuilderAndTags().size(), greaterThan(1));
+            int queryAndTags = queryAndTags(plan);
+            assertThat(queryAndTags, equalTo(points + 1)); // include null bucket
         }
         {
+            int points = between(2, 64);
             String q = String.format(Locale.ROOT, """
                 from test
                 | where date >= "2023-10-19"
                 | stats count(*) by x = round_to(integer, %s)
-                """, points(between(1, 63)));
-            PhysicalPlan plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
-            EsQueryExec esQuery = (EsQueryExec) plan.collectFirstChildren(EsQueryExec.class::isInstance).getFirst();
-            assertThat(esQuery.queryBuilderAndTags().size(), greaterThan(1));
+                """, pointArray(points));
+            var plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
+            int queryAndTags = queryAndTags(plan);
+            assertThat(queryAndTags, equalTo(points + 1)); // include null bucket
         }
         {
+            int points = between(65, 128);
             String q = String.format(Locale.ROOT, """
                 from test
                 | where date >= "2023-10-19"
                 | stats count(*) by x = round_to(integer, %s)
-                """, points(between(65, 128)));
-            PhysicalPlan plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
-            EsQueryExec esQuery = (EsQueryExec) plan.collectFirstChildren(EsQueryExec.class::isInstance).getFirst();
-            assertThat(esQuery.queryBuilderAndTags().size(), equalTo(1));
+                """, pointArray(points));
+            var plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
+            int queryAndTags = queryAndTags(plan);
+            assertThat(queryAndTags, equalTo(1)); // no rewrite
         }
         {
+            int points = between(2, 19);
             String q = String.format(Locale.ROOT, """
                 from test
                 | where date >= "2023-10-19"
                 | where keyword LIKE "w*"
                 | stats count(*) by x = round_to(integer, %s)
-                """, points(between(1, 19)));
-            PhysicalPlan plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
-            EsQueryExec esQuery = (EsQueryExec) plan.collectFirstChildren(EsQueryExec.class::isInstance).getFirst();
-            assertThat(esQuery.queryBuilderAndTags().size(), greaterThan(1));
+                """, pointArray(points));
+            var plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
+            int queryAndTags = queryAndTags(plan);
+            assertThat("points=" + points, queryAndTags, equalTo(points + 1)); // include null bucket
         }
         {
+            int points = between(20, 128);
             String q = String.format(Locale.ROOT, """
                 from test
                 | where date >= "2023-10-19"
                 | where keyword LIKE "*w*"
                 | stats count(*) by x = round_to(integer, %s)
-                """, points(between(20, 128)));
+                """, pointArray(points));
             PhysicalPlan plan = plannerOptimizer.plan(q, searchStats, makeAnalyzer("mapping-all-types.json"));
-            EsQueryExec esQuery = (EsQueryExec) plan.collectFirstChildren(EsQueryExec.class::isInstance).getFirst();
-            assertThat(esQuery.queryBuilderAndTags().size(), equalTo(1));
+            int queryAndTags = queryAndTags(plan);
+            assertThat("points=" + points, queryAndTags, equalTo(1)); // no rewrite
         }
     }