[Spark][#36841] Reactivate structured streaming state tests and prune obsolete DStream helpers

runners/spark/4/src/test/java/org/apache/beam/runners/spark/structuredstreaming/translation/streaming/StreamingPipelineLifecycleTest.java

@@ -33,6 +33,7 @@
 import org.apache.beam.sdk.io.Read;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
+import org.joda.time.Duration;
 import org.junit.ClassRule;
 import org.junit.Rule;
 import org.junit.Test;
@@ -123,6 +124,41 @@ public void cancelStopsTheQueryAndReportsCancelled() throws Exception {
     awaitActiveQueries(count -> count == 0, "the streaming query did not stop after cancel");
   }
 
+  @Test
+  public void timeoutKeepsRunningState() throws Exception {
+    String tag = "lifecycle-timeout";
+    String collectorId = StreamingTestUtils.newCollectorId(tag);
+
+    SparkStructuredStreamingPipelineOptions options =
+        StreamingTestUtils.streamingOptions(checkpointDir);
+    options.setStreamingStopAfterIdleBatches(-1);
+    Pipeline pipeline = Pipeline.create(options);
+
+    pipeline
+        .apply("ReadUnbounded", Read.from(new TestUnboundedSource(tag, 1, 10)))
+        .apply("Collect", ParDo.of(new StreamingTestUtils.CollectDoFn<>(collectorId)));
+
+    PipelineResult result = pipeline.run();
+    try {
+      assertEquals(PipelineResult.State.RUNNING, result.getState());
+
+      awaitActiveQueries(count -> count > 0, "no streaming query started");
+
+      PipelineResult.State stateAfterTimeout = result.waitUntilFinish(Duration.millis(1));
+      assertEquals(PipelineResult.State.RUNNING, stateAfterTimeout);
+      assertEquals(PipelineResult.State.RUNNING, result.getState());
+
+      PipelineResult.State cancelledState = result.cancel();
+      assertEquals(PipelineResult.State.CANCELLED, cancelledState);
+      assertEquals(PipelineResult.State.CANCELLED, result.getState());
+    } finally {
+      if (result.getState() == PipelineResult.State.RUNNING) {
+        result.cancel();
+      }
+      awaitActiveQueries(count -> count == 0, "the streaming query did not stop after cancel");
+    }
+  }
+
   /** A failure in any leaf query surfaces through waitUntilFinish. */
   @Test
   public void failingLeafQueryFailsThePipelineAndStopsHealthySibling() throws Exception {

runners/spark/src/test/java/org/apache/beam/runners/spark/structuredstreaming/StructuredStreamingPipelineStateTest.java

@@ -27,7 +27,6 @@
 import java.io.Serializable;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import org.apache.beam.runners.spark.io.CreateStream;
 import org.apache.beam.runners.spark.structuredstreaming.translation.SparkSessionFactory;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
@@ -36,15 +35,11 @@
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.MapElements;
-import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.SimpleFunction;
-import org.apache.beam.sdk.values.PBegin;
-import org.apache.beam.sdk.values.PCollection;
 import org.apache.spark.TaskContext;
 import org.apache.spark.sql.SparkSession;
 import org.joda.time.Duration;
-import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TestName;
@@ -107,21 +102,6 @@ public void processElement(final ProcessContext c) {
         });
   }
 
-  private PTransform<PBegin, PCollection<String>> getValues(
-      final SparkStructuredStreamingPipelineOptions options) {
-    final boolean doNotSyncWithWatermark = false;
-    return options.isStreaming()
-        ? CreateStream.of(StringUtf8Coder.of(), Duration.millis(1), doNotSyncWithWatermark)
-            .nextBatch("one", "two")
-        : Create.of("one", "two");
-  }
-
-  private SparkStructuredStreamingPipelineOptions getStreamingOptions() {
-    options.setRunner(SparkStructuredStreamingRunner.class);
-    options.setStreaming(true);
-    return options;
-  }
-
   private SparkStructuredStreamingPipelineOptions getBatchOptions() {
     options.setRunner(SparkStructuredStreamingRunner.class);
     options.setStreaming(false); // explicit because options is reused throughout the test.
@@ -131,9 +111,9 @@ private SparkStructuredStreamingPipelineOptions getBatchOptions() {
   private Pipeline getPipeline(final SparkStructuredStreamingPipelineOptions options) {
 
     final Pipeline pipeline = Pipeline.create(options);
-    final String name = testName.getMethodName() + "(isStreaming=" + options.isStreaming() + ")";
+    final String name = testName.getMethodName();
 
-    pipeline.apply(getValues(options)).setCoder(StringUtf8Coder.of()).apply(printParDo(name));
+    pipeline.apply(Create.of("one", "two")).setCoder(StringUtf8Coder.of()).apply(printParDo(name));
 
     return pipeline;
   }
@@ -146,7 +126,7 @@ private void testFailedPipeline(final SparkStructuredStreamingPipelineOptions op
     try {
       final Pipeline pipeline = Pipeline.create(options);
       pipeline
-          .apply(getValues(options))
+          .apply(Create.of("one", "two"))
           .setCoder(StringUtf8Coder.of())
           .apply(
               MapElements.via(
@@ -216,45 +196,21 @@ private void testRunningPipeline(final SparkStructuredStreamingPipelineOptions o
     assertThat(result.waitUntilFinish(), is(PipelineResult.State.CANCELLED));
   }
 
-  @Ignore("TODO: Reactivate with streaming.")
-  @Test
-  public void testStreamingPipelineRunningState() throws Exception {
-    testRunningPipeline(getStreamingOptions());
-  }
-
   @Test
   public void testBatchPipelineRunningState() throws Exception {
     testRunningPipeline(getBatchOptions());
   }
 
-  @Ignore("TODO: Reactivate with streaming.")
-  @Test
-  public void testStreamingPipelineCanceledState() throws Exception {
-    testCanceledPipeline(getStreamingOptions());
-  }
-
   @Test
   public void testBatchPipelineCanceledState() throws Exception {
     testCanceledPipeline(getBatchOptions());
   }
 
-  @Ignore("TODO: Reactivate with streaming.")
-  @Test
-  public void testStreamingPipelineFailedState() throws Exception {
-    testFailedPipeline(getStreamingOptions());
-  }
-
   @Test
   public void testBatchPipelineFailedState() throws Exception {
     testFailedPipeline(getBatchOptions());
   }
 
-  @Ignore("TODO: Reactivate with streaming.")
-  @Test
-  public void testStreamingPipelineTimeoutState() throws Exception {
-    testTimeoutPipeline(getStreamingOptions());
-  }
-
   @Test
   public void testBatchPipelineTimeoutState() throws Exception {
     testTimeoutPipeline(getBatchOptions());