Rename test builder MSQTester.setExpectedSegment (#16837)

* Rename setExpectedSegment to setExpectedSegments in MSQTestBase. * Add expected segments for max num segments test cases.
apache · Aug 2, 2024 · fe6772a · fe6772a
1 parent 9b731e8
commit fe6772a
Show file tree

Hide file tree

Showing 7 changed files with 81 additions and 69 deletions.
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQArraysTest.java
@@ -230,7 +230,7 @@ public void testReplaceMvdWithStringArraySkipValidation(String contextName, Map<
         .setQueryContext(adjustedContext)
         .setExpectedDataSource("foo")
         .setExpectedRowSignature(rowSignature)
-        .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo", Intervals.ETERNITY, "test", 0)))
+        .setExpectedSegments(ImmutableSet.of(SegmentId.of("foo", Intervals.ETERNITY, "test", 0)))
         .setExpectedResultRows(
             NullHandling.sqlCompatible()
             ? ImmutableList.of(
@@ -278,7 +278,7 @@ public void testReplaceMvdWithMvd(String contextName, Map<String, Object> contex
         .setQueryContext(adjustedContext)
         .setExpectedDataSource("foo")
         .setExpectedRowSignature(rowSignature)
-        .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo", Intervals.ETERNITY, "test", 0)))
+        .setExpectedSegments(ImmutableSet.of(SegmentId.of("foo", Intervals.ETERNITY, "test", 0)))
         .setExpectedResultRows(
             ImmutableList.of(
                 new Object[]{0L, null},
@@ -325,7 +325,7 @@ public void testInsertOnFoo1WithMultiValueToArrayGroupByWithDefaultContext(Strin
                      .setExpectedDataSource("foo1")
                      .setExpectedRowSignature(rowSignature)
                      .setQueryContext(context)
-                     .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
+                     .setExpectedSegments(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
                      .setExpectedResultRows(expectedRows)
                      .verifyResults();
   }

diff --git a/...core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQComplexGroupByTest.java b/...core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQComplexGroupByTest.java
@@ -129,7 +129,7 @@ public void testInsertWithoutRollupOnNestedData(String contextName, Map<String,
                              + " GROUP BY 1\n"
                              + " PARTITIONED BY ALL")
                      .setQueryContext(context)
-                     .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
+                     .setExpectedSegments(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
                      .setExpectedDataSource("foo1")
                      .setExpectedRowSignature(RowSignature.builder()
                                                           .add("__time", ColumnType.LONG)
@@ -257,7 +257,7 @@ public void testInsertWithRollupOnNestedData(String contextName, Map<String, Obj
                              + " GROUP BY 1\n"
                              + " PARTITIONED BY ALL")
                      .setQueryContext(adjustedContext)
-                     .setExpectedSegment(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
+                     .setExpectedSegments(ImmutableSet.of(SegmentId.of("foo1", Intervals.ETERNITY, "test", 0)))
                      .setExpectedDataSource("foo1")
                      .setExpectedRowSignature(RowSignature.builder()
                                                           .add("__time", ColumnType.LONG)

diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQExportTest.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/exec/MSQExportTest.java
@@ -57,7 +57,7 @@ public void testExport() throws IOException
                      .setExpectedDataSource("foo1")
                      .setQueryContext(DEFAULT_MSQ_CONTEXT)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();
 
@@ -88,7 +88,7 @@ public void testExport2() throws IOException
                      .setExpectedDataSource("foo1")
                      .setQueryContext(DEFAULT_MSQ_CONTEXT)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();
 
@@ -127,7 +127,7 @@ public void testNumberOfRowsPerFile()
                      .setExpectedDataSource("foo1")
                      .setQueryContext(queryContext)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();
 
@@ -168,7 +168,7 @@ void testExportComplexColumns() throws IOException
                      .setExpectedDataSource("foo1")
                      .setQueryContext(DEFAULT_MSQ_CONTEXT)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();
 
@@ -219,7 +219,7 @@ void testExportSketchColumns() throws IOException
                      .setExpectedDataSource("foo1")
                      .setQueryContext(DEFAULT_MSQ_CONTEXT)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();
 
@@ -258,7 +258,7 @@ void testEmptyExport() throws IOException
                      .setExpectedDataSource("foo1")
                      .setQueryContext(DEFAULT_MSQ_CONTEXT)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();
 
@@ -335,7 +335,7 @@ public void testExportWithLimit() throws IOException
                      .setExpectedDataSource("foo1")
                      .setQueryContext(queryContext)
                      .setExpectedRowSignature(rowSignature)
-                     .setExpectedSegment(ImmutableSet.of())
+                     .setExpectedSegments(ImmutableSet.of())
                      .setExpectedResultRows(ImmutableList.of())
                      .verifyResults();