Simplify CalciteTimechartCommandIT

Signed-off-by: Yuanchun Shen <[email protected]>

Author: yuancu

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalciteTimechartCommandIT.java

@@ -6,6 +6,7 @@
 package org.opensearch.sql.calcite.remote;
 
 import static org.opensearch.sql.legacy.TestUtils.*;
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
 import static org.opensearch.sql.util.MatcherUtils.*;
 
 import java.io.IOException;
@@ -23,6 +24,7 @@ public void init() throws Exception {
     disallowCalciteFallback();
 
     // Create events index with timestamp data
+    loadIndex(Index.BANK);
     loadIndex(Index.EVENTS);
     loadIndex(Index.EVENTS_NULL);
     createEventsManyHostsIndex();
@@ -41,7 +43,6 @@ public void testTimechartWithHourSpanAndGroupBy() throws IOException {
         rows("2024-07-01 00:00:00", "db-01", 1),
         rows("2024-07-01 00:00:00", "web-01", 2),
         rows("2024-07-01 00:00:00", "web-02", 2));
-    assertEquals(3, result.getInt("total"));
   }
 
   @Test
@@ -53,9 +54,6 @@ public void testTimechartWithMinuteSpanAndGroupBy() throws IOException {
         schema("host", "string"),
         schema("count()", "bigint"));
 
-    // Actual result shows 5 rows, not zero-filled results
-    assertEquals(5, result.getInt("total"));
-
     verifyDataRows(
         result,
         rows("2024-07-01 00:00:00", "web-01", 1),
@@ -67,25 +65,13 @@ public void testTimechartWithMinuteSpanAndGroupBy() throws IOException {
 
   @Test
   public void testTimechartWithoutTimestampField() throws IOException {
-    // Create index without @timestamp field
-    String noTimestampMapping =
-        "{\"mappings\":{\"properties\":{\"name\":{\"type\":\"keyword\"},\"occupation\":{\"type\":\"keyword\"},\"country\":{\"type\":\"keyword\"},\"salary\":{\"type\":\"integer\"},\"year\":{\"type\":\"integer\"},\"month\":{\"type\":\"integer\"}}}}";
-    if (!isIndexExist(client(), "no_timestamp")) {
-      createIndexByRestClient(client(), "no_timestamp", noTimestampMapping);
-      loadDataByRestClient(client(), "no_timestamp", "src/test/resources/occupation.json");
-    }
-
-    // Test should throw exception for missing @timestamp field
     Throwable exception =
-        assertThrowsWithReplace(
+        assertThrows(
             ResponseException.class,
             () -> {
-              executeQuery("source=no_timestamp | timechart count()");
+              executeQuery(String.format("source=%s | timechart count()", TEST_INDEX_BANK));
             });
-    assertTrue(
-        "Error message should mention missing @timestamp field",
-        exception.getMessage().contains("@timestamp")
-            || exception.getMessage().contains("timestamp"));
+    verifyErrorMessageContains(exception, "Field [@timestamp] not found.");
   }
 
   @Test
@@ -131,19 +117,8 @@ public void testTimechartWithOtherCategory() throws IOException {
         schema("avg(cpu_usage)", "double"));
 
     // Verify we have 11 data rows (10 hosts + OTHER)
-    assertEquals(11, result.getJSONArray("datarows").length());
-
-    // Verify the OTHER row exists with the correct value
-    boolean foundOther = false;
-    for (int i = 0; i < result.getJSONArray("datarows").length(); i++) {
-      Object[] row = result.getJSONArray("datarows").getJSONArray(i).toList().toArray();
-      if ("OTHER".equals(row[1])) {
-        foundOther = true;
-        assertEquals(35.9, ((Number) row[2]).doubleValue(), 0.01);
-        break;
-      }
-    }
-    assertTrue("OTHER category not found in results", foundOther);
+    verifyNumOfRows(result, 11);
+    verifyDataRowsSome(result, rows("2024-07-01 00:00:00", "OTHER", 35.9));
   }
 
   @Test
@@ -176,9 +151,6 @@ public void testTimechartWithLimitCountGroupBy() throws IOException {
         schema("host", "string"),
         schema("count()", "bigint"));
 
-    // Actual result shows 5 rows, not zero-filled results
-    assertEquals(5, result.getInt("total"));
-
     verifyDataRows(
         result,
         rows("2024-07-01 00:00:00", "web-01", 1),
@@ -202,7 +174,7 @@ public void testTimechartWithLimitZeroAndAvg() throws IOException {
         schema("avg(cpu_usage)", "double"));
 
     // Verify we have 11 data rows (all 11 hosts, no OTHER)
-    assertEquals(11, result.getJSONArray("datarows").length());
+    verifyNumOfRows(result, 11);
 
     // Verify no OTHER category
     boolean foundOther = false;
@@ -228,8 +200,7 @@ public void testTimechartWithLimitZeroAndCount() throws IOException {
         schema("host", "string"),
         schema("count()", "bigint"));
 
-    // For count with limit=0, should show zero-filled results: 11 hosts × 1 time span = 11 rows
-    assertEquals(11, result.getInt("total"));
+    verifyNumOfRows(result, 11);
   }
 
   @Test
@@ -245,7 +216,7 @@ public void testTimechartWithUseOtherFalseAndAvg() throws IOException {
         schema("avg(cpu_usage)", "double"));
 
     // Verify we have 10 data rows (top 10 hosts, no OTHER)
-    assertEquals(10, result.getJSONArray("datarows").length());
+    verifyNumOfRows(result, 10);
 
     // Verify no OTHER category
     boolean foundOther = false;
@@ -270,9 +241,7 @@ public void testTimechartWithUseOtherFalseAndCount() throws IOException {
         schema("host", "string"),
         schema("count()", "bigint"));
 
-    // For count with useother=false, should show zero-filled results: 10 hosts × 1 time span = 10
-    // rows
-    assertEquals(10, result.getInt("total"));
+    verifyNumOfRows(result, 10);
   }
 
   @Test
@@ -285,9 +254,6 @@ public void testTimechartWithCountNoLimitByHostShowZero() throws IOException {
         schema("host", "string"),
         schema("count()", "bigint"));
 
-    // The actual result shows 5 rows, not 15 as zero-filling doesn't happen as expected
-    assertEquals(5, result.getInt("total"));
-
     verifyDataRows(
         result,
         rows("2024-07-01 00:00:00", "web-01", 1),
@@ -310,34 +276,12 @@ public void testTimechartWithLimitAndUseOther() throws IOException {
         schema("host", "string"),
         schema("avg(cpu_usage)", "double"));
 
-    // Verify we have 4 data rows (3 hosts + OTHER)
-    assertEquals(4, result.getJSONArray("datarows").length());
-
-    // Verify specific values with tolerance for floating point precision
-    boolean foundOther = false, foundWeb03 = false, foundWeb07 = false, foundWeb09 = false;
-    for (int i = 0; i < result.getJSONArray("datarows").length(); i++) {
-      Object[] row = result.getJSONArray("datarows").getJSONArray(i).toList().toArray();
-      String host = (String) row[1];
-      double cpuUsage = ((Number) row[2]).doubleValue();
-
-      if ("OTHER".equals(host)) {
-        foundOther = true;
-        assertEquals(41.3, cpuUsage, 0.1);
-      } else if ("web-03".equals(host)) {
-        foundWeb03 = true;
-        assertEquals(55.3, cpuUsage, 0.1);
-      } else if ("web-07".equals(host)) {
-        foundWeb07 = true;
-        assertEquals(48.6, cpuUsage, 0.1);
-      } else if ("web-09".equals(host)) {
-        foundWeb09 = true;
-        assertEquals(67.8, cpuUsage, 0.1);
-      }
-    }
-    assertTrue("OTHER not found", foundOther);
-    assertTrue("web-03 not found", foundWeb03);
-    assertTrue("web-07 not found", foundWeb07);
-    assertTrue("web-09 not found", foundWeb09);
+    verifyDataRows(
+        result,
+        closeTo("2024-07-01 00:00:00", "OTHER", 41.300000000000004),
+        closeTo("2024-07-01 00:00:00", "web-03", 55.3),
+        closeTo("2024-07-01 00:00:00", "web-07", 48.6),
+        closeTo("2024-07-01 00:00:00", "web-09", 67.8));
   }
 
   @Test

integ-test/src/test/java/org/opensearch/sql/util/MatcherUtils.java

@@ -184,6 +184,12 @@ public static void verifyDataRowsInOrder(JSONObject response, Matcher<JSONArray>
     verifyInOrder(response.getJSONArray("datarows"), matchers);
   }
 
+  @SafeVarargs
+  @SuppressWarnings("unchecked")
+  public static void verifyDataRowsSome(JSONObject response, Matcher<JSONArray>... matchers) {
+    verifySome(response.getJSONArray("datarows"), matchers);
+  }
+
   public static void verifyNumOfRows(JSONObject response, int numOfRow) {
     assertEquals(numOfRow, response.getJSONArray("datarows").length());
   }