Extend TestSchedulingUtils with top level union

Author: losipiuk

core/trino-main/src/test/java/io/trino/execution/scheduler/TestSchedulingUtils.java

@@ -72,6 +72,34 @@ public void testCanStreamNoJoin()
         assertThat(SchedulingUtils.canStream(parentSubPlan, valuesSubPlan("a"))).isTrue();
     }
 
+    @Test
+    public void testCanStreamUnion()
+    {
+        /*
+                  parent(union)
+                   /    \    \
+     -------------------------------------- stage boundary
+                  a     b     c
+         */
+        SubPlan aSubPlan = valuesSubPlan("a");
+        RemoteSourceNode remoteSourceA = remoteSource("a");
+
+        SubPlan bSubPlan = valuesSubPlan("b");
+        RemoteSourceNode remoteSourceB = remoteSource("b");
+
+        SubPlan cSubPlan = valuesSubPlan("c");
+        RemoteSourceNode remoteSourceC = remoteSource("c");
+
+        SubPlan parentSubPlan = createSubPlan(
+                "parent",
+                union("union", ImmutableList.of(remoteSourceA, remoteSourceB, remoteSourceC)),
+                ImmutableList.of(aSubPlan, bSubPlan));
+
+        assertThat(SchedulingUtils.canStream(parentSubPlan, aSubPlan)).isTrue();
+        assertThat(SchedulingUtils.canStream(parentSubPlan, bSubPlan)).isTrue();
+        assertThat(SchedulingUtils.canStream(parentSubPlan, cSubPlan)).isTrue();
+    }
+
     @Test
     public void testCanStreamJoin()
     {