content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.beam.runners.spark.translation.streaming;
  
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.allOf;
  import static org.hamcrest.Matchers.equalTo;
  import static org.hamcrest.Matchers.greaterThanOrEqualTo;
  import static org.hamcrest.Matchers.is;
  import static org.hamcrest.Matchers.lessThanOrEqualTo;
  import static org.junit.Assert.assertNotEquals;
  import static org.junit.Assert.fail;
  
  import java.io.IOException;
  import java.io.Serializable;
  import java.util.concurrent.atomic.AtomicInteger;
  import org.apache.beam.runners.spark.SparkPipelineOptions;
  import org.apache.beam.runners.spark.StreamingTest;
  import org.apache.beam.runners.spark.TestSparkPipelineOptions;
  import org.apache.beam.runners.spark.io.CreateStream;
  import org.apache.beam.sdk.coders.KvCoder;
  import org.apache.beam.sdk.coders.StringUtf8Coder;
  import org.apache.beam.sdk.coders.VarIntCoder;
  import org.apache.beam.sdk.coders.VarLongCoder;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.beam.sdk.testing.PAssert;
  import org.apache.beam.sdk.testing.TestPipeline;
  import org.apache.beam.sdk.transforms.Combine;
  import org.apache.beam.sdk.transforms.Count;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.Flatten;
  import org.apache.beam.sdk.transforms.GroupByKey;
  import org.apache.beam.sdk.transforms.ParDo;
  import org.apache.beam.sdk.transforms.SerializableFunction;
  import org.apache.beam.sdk.transforms.Sum;
  import org.apache.beam.sdk.transforms.Values;
  import org.apache.beam.sdk.transforms.WithKeys;
  import org.apache.beam.sdk.transforms.windowing.AfterPane;
  import org.apache.beam.sdk.transforms.windowing.AfterProcessingTime;
  import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  import org.apache.beam.sdk.transforms.windowing.DefaultTrigger;
  import org.apache.beam.sdk.transforms.windowing.FixedWindows;
  import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
  import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
  import org.apache.beam.sdk.transforms.windowing.Never;
  import org.apache.beam.sdk.transforms.windowing.Window;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionList;
  import org.apache.beam.sdk.values.PCollectionTuple;
  import org.apache.beam.sdk.values.TimestampedValue;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.sdk.values.TupleTagList;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.junit.Rule;
  import org.junit.Test;
  import org.junit.experimental.categories.Category;
  import org.junit.rules.ExpectedException;
  
  /**
   * A test suite to test Spark runner implementation of triggers and panes.
   *
   * <p>Since Spark is a micro-batch engine, and will process any test-sized input within the same
   * (first) batch, it is important to make sure inputs are ingested across micro-batches using {@link
   * org.apache.spark.streaming.dstream.QueueInputDStream}. This test suite uses {@link CreateStream}
   * to construct such {@link org.apache.spark.streaming.dstream.QueueInputDStream} and advance the
   * system's WMs. //TODO: add synchronized/processing time trigger.
   */
  @Category(StreamingTest.class)
  public class CreateStreamTest implements Serializable {
  
    @Rule public final transient TestPipeline p = TestPipeline.fromOptions(streamingOptions());
    @Rule public final transient ExpectedException thrown = ExpectedException.none();
  
    @Test
    public void testLateDataAccumulating() throws IOException {
      Instant instant = new Instant(0);
      CreateStream<Integer> source =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(6)))
              .nextBatch(
                  TimestampedValue.of(1, instant),
                  TimestampedValue.of(2, instant),
                  TimestampedValue.of(3, instant))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(20)))
              // These elements are late but within the allowed lateness
              .nextBatch(TimestampedValue.of(4, instant), TimestampedValue.of(5, instant))
              // These elements are droppably late
              .advanceNextBatchWatermarkToInfinity()
              .nextBatch(
                  TimestampedValue.of(-1, instant),
                  TimestampedValue.of(-2, instant),
                  TimestampedValue.of(-3, instant));
  
      PCollection<Integer> windowed =
          p.apply(source)
              .apply(
                  Window.<Integer>into(FixedWindows.of(Duration.standardMinutes(5)))
                      .triggering(
                          AfterWatermark.pastEndOfWindow()
                              .withEarlyFirings(
                                  AfterProcessingTime.pastFirstElementInPane()
                                      .plusDelayOf(Duration.standardMinutes(2)))
                              .withLateFirings(AfterPane.elementCountAtLeast(1)))
                      .accumulatingFiredPanes()
                      .withAllowedLateness(
                          Duration.standardMinutes(5), Window.ClosingBehavior.FIRE_ALWAYS));
      PCollection<Integer> triggered =
          windowed
              .apply(WithKeys.of(1))
              .apply(GroupByKey.create())
              .apply(Values.create())
              .apply(Flatten.iterables());
      PCollection<Long> count =
          windowed.apply(Combine.globally(Count.<Integer>combineFn()).withoutDefaults());
      PCollection<Integer> sum = windowed.apply(Sum.integersGlobally().withoutDefaults());
  
      IntervalWindow window = new IntervalWindow(instant, instant.plus(Duration.standardMinutes(5L)));
      PAssert.that(triggered).inFinalPane(window).containsInAnyOrder(1, 2, 3, 4, 5);
      PAssert.that(triggered).inOnTimePane(window).containsInAnyOrder(1, 2, 3);
      PAssert.that(count)
          .inWindow(window)
          .satisfies(
              input -> {
                for (Long count1 : input) {
                  assertThat(count1, allOf(greaterThanOrEqualTo(3L), lessThanOrEqualTo(5L)));
                }
                return null;
              });
      PAssert.that(sum)
          .inWindow(window)
          .satisfies(
              input -> {
                for (Integer sum1 : input) {
                  assertThat(sum1, allOf(greaterThanOrEqualTo(6), lessThanOrEqualTo(15)));
                }
                return null;
              });
  
      p.run();
    }
  
    @Test
    public void testDiscardingMode() throws IOException {
      CreateStream<String> source =
          CreateStream.of(StringUtf8Coder.of(), batchDuration())
              .nextBatch(
                  TimestampedValue.of("firstPane", new Instant(100)),
                  TimestampedValue.of("alsoFirstPane", new Instant(200)))
              .advanceWatermarkForNextBatch(new Instant(1001L))
              .nextBatch(TimestampedValue.of("onTimePane", new Instant(500)))
              .advanceNextBatchWatermarkToInfinity()
              .nextBatch(
                  TimestampedValue.of("finalLatePane", new Instant(750)),
                  TimestampedValue.of("alsoFinalLatePane", new Instant(250)));
  
      FixedWindows windowFn = FixedWindows.of(Duration.millis(1000L));
      Duration allowedLateness = Duration.millis(5000L);
      PCollection<String> values =
          p.apply(source)
              .apply(
                  Window.<String>into(windowFn)
                      .triggering(
                          AfterWatermark.pastEndOfWindow()
                              .withEarlyFirings(AfterPane.elementCountAtLeast(2))
                              .withLateFirings(Never.ever()))
                      .discardingFiredPanes()
                      .withAllowedLateness(allowedLateness))
              .apply(WithKeys.of(1))
              .apply(GroupByKey.create())
              .apply(Values.create())
              .apply(Flatten.iterables());
  
      IntervalWindow window = windowFn.assignWindow(new Instant(100));
      PAssert.that(values)
          .inWindow(window)
          .containsInAnyOrder(
              "firstPane", "alsoFirstPane", "onTimePane", "finalLatePane", "alsoFinalLatePane");
      PAssert.that(values)
          .inCombinedNonLatePanes(window)
          .containsInAnyOrder("firstPane", "alsoFirstPane", "onTimePane");
      PAssert.that(values).inOnTimePane(window).containsInAnyOrder("onTimePane");
      PAssert.that(values)
          .inFinalPane(window)
          .containsInAnyOrder("finalLatePane", "alsoFinalLatePane");
  
      p.run();
    }
  
    @Test
    public void testFirstElementLate() throws IOException {
      Instant lateElementTimestamp = new Instant(-1_000_000);
      CreateStream<String> source =
          CreateStream.of(StringUtf8Coder.of(), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(new Instant(0))
              .emptyBatch()
              .nextBatch(
                  TimestampedValue.of("late", lateElementTimestamp),
                  TimestampedValue.of("onTime", new Instant(100)))
              .advanceNextBatchWatermarkToInfinity();
  
      FixedWindows windowFn = FixedWindows.of(Duration.millis(1000L));
      Duration allowedLateness = Duration.millis(5000L);
      PCollection<String> values =
          p.apply(source)
              .apply(
                  Window.<String>into(windowFn)
                      .triggering(DefaultTrigger.of())
                      .discardingFiredPanes()
                      .withAllowedLateness(allowedLateness))
              .apply(WithKeys.of(1))
              .apply(GroupByKey.create())
              .apply(Values.create())
              .apply(Flatten.iterables());
  
      PAssert.that(values).inWindow(windowFn.assignWindow(lateElementTimestamp)).empty();
      PAssert.that(values)
          .inWindow(windowFn.assignWindow(new Instant(100)))
          .containsInAnyOrder("onTime");
  
      p.run();
    }
  
    @Test
    public void testElementsAtAlmostPositiveInfinity() throws IOException {
      Instant endOfGlobalWindow = GlobalWindow.INSTANCE.maxTimestamp();
      CreateStream<String> source =
          CreateStream.of(StringUtf8Coder.of(), batchDuration())
              .nextBatch(
                  TimestampedValue.of("foo", endOfGlobalWindow),
                  TimestampedValue.of("bar", endOfGlobalWindow))
              .advanceNextBatchWatermarkToInfinity();
  
      FixedWindows windows = FixedWindows.of(Duration.standardHours(6));
      PCollection<String> windowedValues =
          p.apply(source)
              .apply(Window.into(windows))
              .apply(WithKeys.of(1))
              .apply(GroupByKey.create())
              .apply(Values.create())
              .apply(Flatten.iterables());
  
      PAssert.that(windowedValues)
          .inWindow(windows.assignWindow(GlobalWindow.INSTANCE.maxTimestamp()))
          .containsInAnyOrder("foo", "bar");
      p.run();
    }
  
    @Test
    public void testMultipleStreams() throws IOException {
      CreateStream<String> source =
          CreateStream.of(StringUtf8Coder.of(), batchDuration())
              .nextBatch("foo", "bar")
              .advanceNextBatchWatermarkToInfinity();
      CreateStream<Integer> other =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .nextBatch(1, 2, 3, 4)
              .advanceNextBatchWatermarkToInfinity();
  
      PCollection<String> createStrings =
          p.apply("CreateStrings", source)
              .apply(
                  "WindowStrings",
                  Window.<String>configure()
                      .triggering(AfterPane.elementCountAtLeast(2))
                      .withAllowedLateness(Duration.ZERO)
                      .accumulatingFiredPanes());
      PAssert.that(createStrings).containsInAnyOrder("foo", "bar");
  
      PCollection<Integer> createInts =
          p.apply("CreateInts", other)
              .apply(
                  "WindowInts",
                  Window.<Integer>configure()
                      .triggering(AfterPane.elementCountAtLeast(4))
                      .withAllowedLateness(Duration.ZERO)
                      .accumulatingFiredPanes());
      PAssert.that(createInts).containsInAnyOrder(1, 2, 3, 4);
  
      p.run();
    }
  
    @Test
    public void testFlattenedWithWatermarkHold() throws IOException {
      Instant instant = new Instant(0);
      CreateStream<Integer> source1 =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(5)))
              .nextBatch(
                  TimestampedValue.of(1, instant),
                  TimestampedValue.of(2, instant),
                  TimestampedValue.of(3, instant))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(10)));
      CreateStream<Integer> source2 =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(1)))
              .nextBatch(TimestampedValue.of(4, instant))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(2)))
              .nextBatch(TimestampedValue.of(5, instant))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(5)))
              .emptyBatch()
              .advanceNextBatchWatermarkToInfinity();
  
      PCollection<Integer> windowed1 =
          p.apply("CreateStream1", source1)
              .apply(
                  "Window1",
                  Window.<Integer>into(FixedWindows.of(Duration.standardMinutes(5)))
                      .triggering(AfterWatermark.pastEndOfWindow())
                      .accumulatingFiredPanes()
                      .withAllowedLateness(Duration.ZERO));
      PCollection<Integer> windowed2 =
          p.apply("CreateStream2", source2)
              .apply(
                  "Window2",
                  Window.<Integer>into(FixedWindows.of(Duration.standardMinutes(5)))
                      .triggering(AfterWatermark.pastEndOfWindow())
                      .accumulatingFiredPanes()
                      .withAllowedLateness(Duration.ZERO));
  
      PCollectionList<Integer> pCollectionList = PCollectionList.of(windowed1).and(windowed2);
      PCollection<Integer> flattened = pCollectionList.apply(Flatten.pCollections());
      PCollection<Integer> triggered =
          flattened
              .apply(WithKeys.of(1))
              .apply(GroupByKey.create())
              .apply(Values.create())
              .apply(Flatten.iterables());
  
      IntervalWindow window = new IntervalWindow(instant, instant.plus(Duration.standardMinutes(5L)));
      PAssert.that(triggered).inOnTimePane(window).containsInAnyOrder(1, 2, 3, 4, 5);
  
      p.run();
    }
  
    /**
     * Test multiple output {@link ParDo} in streaming pipelines. This is currently needed as a test
     * for https://issues.apache.org/jira/browse/BEAM-2029 since {@link
     * org.apache.beam.sdk.testing.ValidatesRunner} tests do not currently run for Spark runner in
     * streaming mode.
     */
    @Test
    public void testMultiOutputParDo() throws IOException {
      Instant instant = new Instant(0);
      CreateStream<Integer> source1 =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(5)))
              .nextBatch(
                  TimestampedValue.of(1, instant),
                  TimestampedValue.of(2, instant),
                  TimestampedValue.of(3, instant))
              .advanceNextBatchWatermarkToInfinity();
  
      PCollection<Integer> inputs = p.apply(source1);
  
      final TupleTag<Integer> mainTag = new TupleTag<>();
      final TupleTag<Integer> additionalTag = new TupleTag<>();
  
      PCollectionTuple outputs =
          inputs.apply(
              ParDo.of(
                      new DoFn<Integer, Integer>() {
  
                        @SuppressWarnings("unused")
                        @ProcessElement
                        public void process(ProcessContext context) {
                          Integer element = context.element();
                          context.output(element);
                          context.output(additionalTag, element + 1);
                        }
                      })
                  .withOutputTags(mainTag, TupleTagList.of(additionalTag)));
  
      PCollection<Integer> output1 = outputs.get(mainTag).setCoder(VarIntCoder.of());
      PCollection<Integer> output2 = outputs.get(additionalTag).setCoder(VarIntCoder.of());
  
      PAssert.that(output1).containsInAnyOrder(1, 2, 3);
      PAssert.that(output2).containsInAnyOrder(2, 3, 4);
  
      p.run();
    }
  
    /**
     * Test that {@link ParDo} aligns both setup and teardown calls in streaming pipelines. See
     * https://issues.apache.org/jira/browse/BEAM-6859.
     */
    @Test
    public void testParDoCallsSetupAndTeardown() {
      Instant instant = new Instant(0);
  
      p.apply(
              CreateStream.of(VarIntCoder.of(), batchDuration())
                  .emptyBatch()
                  .advanceWatermarkForNextBatch(instant.plus(Duration.standardMinutes(5)))
                  .nextBatch(
                      TimestampedValue.of(1, instant),
                      TimestampedValue.of(2, instant),
                      TimestampedValue.of(3, instant))
                  .advanceNextBatchWatermarkToInfinity())
          .apply(ParDo.of(new LifecycleDoFn()));
  
      p.run();
  
      assertThat(
          "Function should have been torn down",
          LifecycleDoFn.teardownCalls.intValue(),
          is(equalTo(LifecycleDoFn.setupCalls.intValue())));
    }
  
    @Test
    public void testElementAtPositiveInfinityThrows() {
      CreateStream<Integer> source =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .nextBatch(
                  TimestampedValue.of(
                      -1, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1L))))
              .advanceNextBatchWatermarkToInfinity();
      thrown.expect(IllegalArgumentException.class);
      source.nextBatch(TimestampedValue.of(1, BoundedWindow.TIMESTAMP_MAX_VALUE));
    }
  
    @Test
    public void testAdvanceWatermarkNonMonotonicThrows() {
      CreateStream<Integer> source =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .advanceWatermarkForNextBatch(new Instant(0L));
      thrown.expect(IllegalArgumentException.class);
      source.advanceWatermarkForNextBatch(new Instant(-1L)).advanceNextBatchWatermarkToInfinity();
    }
  
    @Test
    public void testAdvanceWatermarkEqualToPositiveInfinityThrows() {
      CreateStream<Integer> source =
          CreateStream.of(VarIntCoder.of(), batchDuration())
              .advanceWatermarkForNextBatch(
                  BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.millis(1L)));
      thrown.expect(IllegalArgumentException.class);
      source.advanceWatermarkForNextBatch(BoundedWindow.TIMESTAMP_MAX_VALUE);
    }
  
    @Test
    public void testInStreamingModeCountByKey() throws Exception {
      Instant instant = new Instant(0);
  
      CreateStream<KV<Integer, Long>> kvSource =
          CreateStream.of(KvCoder.of(VarIntCoder.of(), VarLongCoder.of()), batchDuration())
              .emptyBatch()
              .advanceWatermarkForNextBatch(instant)
              .nextBatch(
                  TimestampedValue.of(KV.of(1, 100L), instant.plus(Duration.standardSeconds(3L))),
                  TimestampedValue.of(KV.of(1, 300L), instant.plus(Duration.standardSeconds(4L))))
              .advanceWatermarkForNextBatch(instant.plus(Duration.standardSeconds(7L)))
              .nextBatch(
                  TimestampedValue.of(KV.of(1, 400L), instant.plus(Duration.standardSeconds(8L))))
              .advanceNextBatchWatermarkToInfinity();
  
      PCollection<KV<Integer, Long>> output =
          p.apply("create kv Source", kvSource)
              .apply(
                  "window input",
                  Window.<KV<Integer, Long>>into(FixedWindows.of(Duration.standardSeconds(3L)))
                      .withAllowedLateness(Duration.ZERO))
              .apply(Count.perKey());
  
      PAssert.that("Wrong count value ", output)
          .satisfies(
              (SerializableFunction<Iterable<KV<Integer, Long>>, Void>)
                  input -> {
                    for (KV<Integer, Long> element : input) {
                      if (element.getKey() == 1) {
                        Long countValue = element.getValue();
                        assertNotEquals("Count Value is 0 !!!", 0L, countValue.longValue());
                      } else {
                        fail("Unknown key in the output PCollection");
                      }
                    }
                    return null;
                  });
      p.run();
    }
  
    private Duration batchDuration() {
      return Duration.millis(p.getOptions().as(SparkPipelineOptions.class).getBatchIntervalMillis());
    }
  
    private static class LifecycleDoFn extends DoFn<Integer, Integer> {
      static AtomicInteger setupCalls = new AtomicInteger(0);
      static AtomicInteger teardownCalls = new AtomicInteger(0);
  
      @Setup
      public void setup() {
        setupCalls.incrementAndGet();
      }
  
      @Teardown
      public void teardown() {
        teardownCalls.incrementAndGet();
      }
  
      @SuppressWarnings("unused")
      @ProcessElement
      public void process(ProcessContext context) {
        Integer element = context.element();
        context.output(element);
      }
    }
  
    static PipelineOptions streamingOptions() {
      PipelineOptions options = TestPipeline.testingPipelineOptions();
      options.as(TestSparkPipelineOptions.class).setStreaming(true);
      return options;
    }
  }
  