Esql refactor date tests (#117923)

This refactors a bit of the type logic in the parametrized testing to pass the input values as java Instants for millisecond and nanosecond date. Mainly, this impacts verifier functions. The goal here is to ensure that the values are correctly converted based on the type they were generated as, rather than relying on the verifier function to know how to convert from a long with no additional information. This will make tests that have mixed millisecond and nanosecond inputs easier to write correctly.

Author: not-napoleon

x-pack/plugin/esql-core/src/main/java/org/elasticsearch/xpack/esql/core/util/DateUtils.java

@@ -174,6 +174,10 @@ public static ZonedDateTime asDateTime(long millis) {
         return ZonedDateTime.ofInstant(Instant.ofEpochMilli(millis), UTC);
     }
 
+    public static ZonedDateTime asDateTime(Instant instant) {
+        return ZonedDateTime.ofInstant(instant, UTC);
+    }
+
     public static long asMillis(ZonedDateTime zonedDateTime) {
         return zonedDateTime.toInstant().toEpochMilli();
     }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/TestCaseSupplier.java

@@ -620,70 +620,6 @@ public static void forUnaryBoolean(
         unary(suppliers, expectedEvaluatorToString, booleanCases(), expectedType, v -> expectedValue.apply((Boolean) v), warnings);
     }
 
-    /**
-     * Generate positive test cases for a unary function operating on an {@link DataType#DATETIME}.
-     * This variant defaults to maximum range of possible values
-     */
-    public static void forUnaryDatetime(
-        List<TestCaseSupplier> suppliers,
-        String expectedEvaluatorToString,
-        DataType expectedType,
-        Function<Instant, Object> expectedValue,
-        List<String> warnings
-    ) {
-        unaryNumeric(
-            suppliers,
-            expectedEvaluatorToString,
-            dateCases(),
-            expectedType,
-            n -> expectedValue.apply(Instant.ofEpochMilli(n.longValue())),
-            warnings
-        );
-    }
-
-    /**
-     * Generate positive test cases for a unary function operating on an {@link DataType#DATETIME}.
-     * This variant accepts a range of values
-     */
-    public static void forUnaryDatetime(
-        List<TestCaseSupplier> suppliers,
-        String expectedEvaluatorToString,
-        DataType expectedType,
-        long min,
-        long max,
-        Function<Instant, Object> expectedValue,
-        List<String> warnings
-    ) {
-        unaryNumeric(
-            suppliers,
-            expectedEvaluatorToString,
-            dateCases(min, max),
-            expectedType,
-            n -> expectedValue.apply(Instant.ofEpochMilli(n.longValue())),
-            warnings
-        );
-    }
-
-    /**
-     * Generate positive test cases for a unary function operating on an {@link DataType#DATE_NANOS}.
-     */
-    public static void forUnaryDateNanos(
-        List<TestCaseSupplier> suppliers,
-        String expectedEvaluatorToString,
-        DataType expectedType,
-        Function<Instant, Object> expectedValue,
-        List<String> warnings
-    ) {
-        unaryNumeric(
-            suppliers,
-            expectedEvaluatorToString,
-            dateNanosCases(),
-            expectedType,
-            n -> expectedValue.apply(DateUtils.toInstant((long) n)),
-            warnings
-        );
-    }
-
     /**
      * Generate positive test cases for a unary function operating on an {@link DataType#GEO_POINT}.
      */
@@ -1912,11 +1848,19 @@ public List<Object> multiRowData() {
         }
 
         /**
-         * @return the data value being supplied, casting unsigned longs into BigIntegers correctly
+         * @return the data value being supplied, casting to java objects when appropriate
          */
         public Object getValue() {
-            if (type == DataType.UNSIGNED_LONG && data instanceof Long l) {
-                return NumericUtils.unsignedLongAsBigInteger(l);
+            if (data instanceof Long l) {
+                if (type == DataType.UNSIGNED_LONG) {
+                    return NumericUtils.unsignedLongAsBigInteger(l);
+                }
+                if (type == DataType.DATETIME) {
+                    return Instant.ofEpochMilli(l);
+                }
+                if (type == DataType.DATE_NANOS) {
+                    return DateUtils.toInstant(l);
+                }
             }
             return data;
         }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToDateNanosTests.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
 
 import java.math.BigInteger;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Supplier;
@@ -36,14 +37,20 @@ public static Iterable<Object[]> parameters() {
         final String read = "Attribute[channel=0]";
         final List<TestCaseSupplier> suppliers = new ArrayList<>();
 
-        TestCaseSupplier.forUnaryDateNanos(suppliers, read, DataType.DATE_NANOS, DateUtils::toLong, List.of());
-        TestCaseSupplier.forUnaryDatetime(
+        TestCaseSupplier.unary(
+            suppliers,
+            read,
+            TestCaseSupplier.dateNanosCases(),
+            DataType.DATE_NANOS,
+            v -> DateUtils.toLong((Instant) v),
+            List.of()
+        );
+        TestCaseSupplier.unary(
             suppliers,
             "ToDateNanosFromDatetimeEvaluator[field=" + read + "]",
+            TestCaseSupplier.dateCases(0, DateUtils.MAX_NANOSECOND_INSTANT.toEpochMilli()),
             DataType.DATE_NANOS,
-            0,
-            DateUtils.MAX_NANOSECOND_INSTANT.toEpochMilli(),
-            i -> DateUtils.toNanoSeconds(i.toEpochMilli()),
+            i -> DateUtils.toNanoSeconds(((Instant) i).toEpochMilli()),
             List.of()
         );
         TestCaseSupplier.forUnaryLong(

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToDatetimeTests.java

@@ -37,12 +37,20 @@ public static Iterable<Object[]> parameters() {
         final String read = "Attribute[channel=0]";
         final List<TestCaseSupplier> suppliers = new ArrayList<>();
 
-        TestCaseSupplier.forUnaryDatetime(suppliers, read, DataType.DATETIME, Instant::toEpochMilli, emptyList());
-        TestCaseSupplier.forUnaryDateNanos(
+        TestCaseSupplier.unary(
+            suppliers,
+            read,
+            TestCaseSupplier.dateCases(),
+            DataType.DATETIME,
+            v -> ((Instant) v).toEpochMilli(),
+            emptyList()
+        );
+        TestCaseSupplier.unary(
             suppliers,
             "ToDatetimeFromDateNanosEvaluator[field=" + read + "]",
+            TestCaseSupplier.dateNanosCases(),
             DataType.DATETIME,
-            i -> DateUtils.toMilliSeconds(DateUtils.toLong(i)),
+            i -> DateUtils.toMilliSeconds(DateUtils.toLong((Instant) i)),
             emptyList()
         );
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToDoubleTests.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
 
 import java.math.BigInteger;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
@@ -49,11 +50,12 @@ public static Iterable<Object[]> parameters() {
         );
 
         TestCaseSupplier.forUnaryBoolean(suppliers, evaluatorName.apply("Boolean"), DataType.DOUBLE, b -> b ? 1d : 0d, List.of());
-        TestCaseSupplier.forUnaryDatetime(
+        TestCaseSupplier.unary(
             suppliers,
             evaluatorName.apply("Long"),
+            TestCaseSupplier.dateCases(),
             DataType.DOUBLE,
-            i -> (double) i.toEpochMilli(),
+            i -> (double) ((Instant) i).toEpochMilli(),
             List.of()
         );
         // random strings that don't look like a double

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToIntegerTests.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
 
 import java.math.BigInteger;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
@@ -48,7 +49,7 @@ public static Iterable<Object[]> parameters() {
             evaluatorName.apply("Long"),
             dateCases(0, Integer.MAX_VALUE),
             DataType.INTEGER,
-            l -> ((Long) l).intValue(),
+            l -> Long.valueOf(((Instant) l).toEpochMilli()).intValue(),
             List.of()
         );
         // datetimes that fall outside Integer's range
@@ -60,7 +61,9 @@ public static Iterable<Object[]> parameters() {
             l -> null,
             l -> List.of(
                 "Line -1:-1: evaluation of [] failed, treating result as null. Only first 20 failures recorded.",
-                "Line -1:-1: org.elasticsearch.xpack.esql.core.InvalidArgumentException: [" + l + "] out of [integer] range"
+                "Line -1:-1: org.elasticsearch.xpack.esql.core.InvalidArgumentException: ["
+                    + ((Instant) l).toEpochMilli()
+                    + "] out of [integer] range"
             )
         );
         // random strings that don't look like an Integer

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToLongTests.java

@@ -43,8 +43,15 @@ public static Iterable<Object[]> parameters() {
         TestCaseSupplier.forUnaryBoolean(suppliers, evaluatorName.apply("Boolean"), DataType.LONG, b -> b ? 1L : 0L, List.of());
 
         // datetimes
-        TestCaseSupplier.forUnaryDatetime(suppliers, read, DataType.LONG, Instant::toEpochMilli, List.of());
-        TestCaseSupplier.forUnaryDateNanos(suppliers, read, DataType.LONG, DateUtils::toLong, List.of());
+        TestCaseSupplier.unary(suppliers, read, TestCaseSupplier.dateCases(), DataType.LONG, v -> ((Instant) v).toEpochMilli(), List.of());
+        TestCaseSupplier.unary(
+            suppliers,
+            read,
+            TestCaseSupplier.dateNanosCases(),
+            DataType.LONG,
+            v -> DateUtils.toLong((Instant) v),
+            List.of()
+        );
         // random strings that don't look like a long
         TestCaseSupplier.forUnaryStrings(
             suppliers,

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToStringTests.java

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
 
 import java.math.BigInteger;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Supplier;
@@ -81,18 +82,20 @@ public static Iterable<Object[]> parameters() {
             b -> new BytesRef(b.toString()),
             List.of()
         );
-        TestCaseSupplier.forUnaryDatetime(
+        TestCaseSupplier.unary(
             suppliers,
             "ToStringFromDatetimeEvaluator[field=" + read + "]",
+            TestCaseSupplier.dateCases(),
             DataType.KEYWORD,
-            i -> new BytesRef(DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER.formatMillis(i.toEpochMilli())),
+            i -> new BytesRef(DateFieldMapper.DEFAULT_DATE_TIME_FORMATTER.formatMillis(((Instant) i).toEpochMilli())),
             List.of()
         );
-        TestCaseSupplier.forUnaryDateNanos(
+        TestCaseSupplier.unary(
             suppliers,
             "ToStringFromDateNanosEvaluator[field=" + read + "]",
+            TestCaseSupplier.dateNanosCases(),
             DataType.KEYWORD,
-            i -> new BytesRef(DateFieldMapper.DEFAULT_DATE_TIME_NANOS_FORMATTER.formatNanos(DateUtils.toLong(i))),
+            i -> new BytesRef(DateFieldMapper.DEFAULT_DATE_TIME_NANOS_FORMATTER.formatNanos(DateUtils.toLong((Instant) i))),
             List.of()
         );
         TestCaseSupplier.forUnaryGeoPoint(

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToUnsignedLongTests.java

@@ -19,6 +19,7 @@
 
 import java.math.BigDecimal;
 import java.math.BigInteger;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
@@ -58,11 +59,12 @@ public static Iterable<Object[]> parameters() {
         );
 
         // datetimes
-        TestCaseSupplier.forUnaryDatetime(
+        TestCaseSupplier.unary(
             suppliers,
             evaluatorName.apply("Long"),
+            TestCaseSupplier.dateCases(),
             DataType.UNSIGNED_LONG,
-            instant -> BigInteger.valueOf(instant.toEpochMilli()),
+            instant -> BigInteger.valueOf(((Instant) instant).toEpochMilli()),
             List.of()
         );
         // random strings that don't look like an unsigned_long

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/predicate/operator/arithmetic/AddTests.java

@@ -159,7 +159,7 @@ public static Iterable<Object[]> parameters() {
         };
         BiFunction<TestCaseSupplier.TypedData, TestCaseSupplier.TypedData, List<String>> warnings = (lhs, rhs) -> {
             try {
-                addDatesAndTemporalAmount(lhs.data(), rhs.data(), AddTests::addMillis);
+                addDatesAndTemporalAmount(lhs.getValue(), rhs.getValue(), AddTests::addMillis);
                 return List.of();
             } catch (ArithmeticException e) {
                 return List.of(
@@ -193,6 +193,7 @@ public static Iterable<Object[]> parameters() {
 
         BinaryOperator<Object> nanosResult = (lhs, rhs) -> {
             try {
+                assert (lhs instanceof Instant) || (rhs instanceof Instant);
                 return addDatesAndTemporalAmount(lhs, rhs, AddTests::addNanos);
             } catch (ArithmeticException e) {
                 return null;
@@ -327,29 +328,28 @@ private static String addErrorMessageString(boolean includeOrdinal, List<Set<Dat
         }
     }
 
-    private static Object addDatesAndTemporalAmount(Object lhs, Object rhs, ToLongBiFunction<Long, TemporalAmount> adder) {
+    private static Object addDatesAndTemporalAmount(Object lhs, Object rhs, ToLongBiFunction<Instant, TemporalAmount> adder) {
         // this weird casting dance makes the expected value lambda symmetric
-        Long date;
+        Instant date;
         TemporalAmount period;
-        if (lhs instanceof Long) {
-            date = (Long) lhs;
+        assert (lhs instanceof Instant) || (rhs instanceof Instant);
+        if (lhs instanceof Instant) {
+            date = (Instant) lhs;
             period = (TemporalAmount) rhs;
         } else {
-            date = (Long) rhs;
+            date = (Instant) rhs;
             period = (TemporalAmount) lhs;
         }
         return adder.applyAsLong(date, period);
     }
 
-    private static long addMillis(Long date, TemporalAmount period) {
+    private static long addMillis(Instant date, TemporalAmount period) {
         return asMillis(asDateTime(date).plus(period));
     }
 
-    private static long addNanos(Long date, TemporalAmount period) {
+    private static long addNanos(Instant date, TemporalAmount period) {
         return DateUtils.toLong(
-            Instant.from(
-                ZonedDateTime.ofInstant(DateUtils.toInstant(date), org.elasticsearch.xpack.esql.core.util.DateUtils.UTC).plus(period)
-            )
+            Instant.from(ZonedDateTime.ofInstant(date, org.elasticsearch.xpack.esql.core.util.DateUtils.UTC).plus(period))
         );
     }