fix coverage

Signed-off-by: Lantao Jin <[email protected]>
opensearch-project · Oct 25, 2024 · df64ccd · df64ccd
1 parent 549bd1a
commit df64ccd
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 0 deletions.
diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/collector/RoundingTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/collector/RoundingTest.java
@@ -5,14 +5,18 @@
 
 package org.opensearch.sql.planner.physical.collector;
 
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.opensearch.sql.data.type.ExprCoreType.DATE;
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 import static org.opensearch.sql.data.type.ExprCoreType.TIME;
+import static org.opensearch.sql.data.type.ExprCoreType.TIMESTAMP;
 
 import org.junit.jupiter.api.Test;
 import org.opensearch.sql.data.model.ExprTimeValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
+import org.opensearch.sql.data.type.ExprType;
 import org.opensearch.sql.exception.ExpressionEvaluationException;
 import org.opensearch.sql.expression.DSL;
 import org.opensearch.sql.expression.span.SpanExpression;
@@ -26,6 +30,35 @@ void time_rounding_illegal_span() {
         ExpressionEvaluationException.class, () -> rounding.round(new ExprTimeValue("23:30:00")));
   }
 
+  @Test
+  void datetime_rounding_span() {
+    SpanExpression dateSpan = DSL.span(DSL.ref("date", DATE), DSL.literal(1), "d");
+    Rounding rounding = Rounding.createRounding(dateSpan);
+    assertInstanceOf(Rounding.DateRounding.class, rounding);
+    SpanExpression timeSpan = DSL.span(DSL.ref("time", TIME), DSL.literal(1), "h");
+    rounding = Rounding.createRounding(timeSpan);
+    assertInstanceOf(Rounding.TimeRounding.class, rounding);
+    SpanExpression timestampSpan = DSL.span(DSL.ref("timestamp", TIMESTAMP), DSL.literal(1), "h");
+    rounding = Rounding.createRounding(timestampSpan);
+    assertInstanceOf(Rounding.TimestampRounding.class, rounding);
+  }
+
+  @Test
+  void datetime_rounding_non_core_type_span() {
+    SpanExpression dateSpan =
+        DSL.span(DSL.ref("date", new MockDateExprType()), DSL.literal(1), "d");
+    Rounding rounding = Rounding.createRounding(dateSpan);
+    assertInstanceOf(Rounding.DateRounding.class, rounding);
+    SpanExpression timeSpan =
+        DSL.span(DSL.ref("time", new MockTimeExprType()), DSL.literal(1), "h");
+    rounding = Rounding.createRounding(timeSpan);
+    assertInstanceOf(Rounding.TimeRounding.class, rounding);
+    SpanExpression timestampSpan =
+        DSL.span(DSL.ref("timestamp", new MockTimestampExprType()), DSL.literal(1), "h");
+    rounding = Rounding.createRounding(timestampSpan);
+    assertInstanceOf(Rounding.TimestampRounding.class, rounding);
+  }
+
   @Test
   void round_unknown_type() {
     SpanExpression span = DSL.span(DSL.ref("unknown", STRING), DSL.literal(1), "");
@@ -41,4 +74,25 @@ void resolve() {
         () -> Rounding.DateTimeUnit.resolve(illegalUnit),
         "Unable to resolve unit " + illegalUnit);
   }
+
+  static class MockDateExprType implements ExprType {
+    @Override
+    public String typeName() {
+      return "DATE";
+    }
+  }
+
+  static class MockTimeExprType implements ExprType {
+    @Override
+    public String typeName() {
+      return "TIME";
+    }
+  }
+
+  static class MockTimestampExprType implements ExprType {
+    @Override
+    public String typeName() {
+      return "TIMESTAMP";
+    }
+  }
 }
diff --git a/integ-test/src/test/java/org/opensearch/sql/ppl/DateTimeImplementationIT.java b/integ-test/src/test/java/org/opensearch/sql/ppl/DateTimeImplementationIT.java
@@ -190,4 +190,27 @@ public void testSpanDatetimeWithCustomFormat() throws IOException {
     verifySchema(result, schema("cnt", null, "integer"), schema("span", null, "date"));
     verifyDataRows(result, rows(2, "1984-04-12"));
   }
+
+  @Test
+  public void testSpanDatetimeWithEpochMillisFormat() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | eval a = 1 | stats count() as cnt by span(epoch_millis, 1d) as span",
+                TEST_INDEX_DATE_FORMATS));
+    verifySchema(result, schema("cnt", null, "integer"), schema("span", null, "timestamp"));
+    verifyDataRows(result, rows(2, "1984-04-12 00:00:00"));
+  }
+
+  @Test
+  public void testSpanDatetimeWithDisjunctiveDifferentFormats() throws IOException {
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | eval a = 1 | stats count() as cnt by span(yyyy-MM-dd_OR_epoch_millis,"
+                    + " 1d) as span",
+                TEST_INDEX_DATE_FORMATS));
+    verifySchema(result, schema("cnt", null, "integer"), schema("span", null, "timestamp"));
+    verifyDataRows(result, rows(2, "1984-04-12 00:00:00"));
+  }
 }