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.flink.translation.wrappers.streaming;
  
  import static org.apache.beam.runners.flink.translation.wrappers.streaming.StreamRecordStripper.stripStreamRecordFromWindowedValue;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.containsInAnyOrder;
  import static org.hamcrest.Matchers.emptyIterable;
  import static org.hamcrest.Matchers.equalTo;
  import static org.hamcrest.Matchers.greaterThan;
  import static org.hamcrest.Matchers.instanceOf;
  import static org.hamcrest.Matchers.is;
  import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertThrows;
  import static org.junit.Assert.assertTrue;
  
  import com.fasterxml.jackson.databind.type.TypeFactory;
  import com.fasterxml.jackson.databind.util.LRUMap;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.function.Supplier;
  import java.util.stream.Collectors;
  import org.apache.beam.runners.core.DoFnRunner;
  import org.apache.beam.runners.core.StateNamespace;
  import org.apache.beam.runners.core.StateNamespaces;
  import org.apache.beam.runners.core.StateTag;
  import org.apache.beam.runners.core.StateTags;
  import org.apache.beam.runners.core.StatefulDoFnRunner;
  import org.apache.beam.runners.core.StepContext;
  import org.apache.beam.runners.core.TimerInternals;
  import org.apache.beam.runners.core.construction.SerializablePipelineOptions;
  import org.apache.beam.runners.flink.FlinkPipelineOptions;
  import org.apache.beam.runners.flink.adapter.FlinkKey;
  import org.apache.beam.runners.flink.metrics.FlinkMetricContainer;
  import org.apache.beam.runners.flink.translation.types.CoderTypeSerializer;
  import org.apache.beam.sdk.Pipeline;
  import org.apache.beam.sdk.coders.Coder;
  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.state.StateSpec;
  import org.apache.beam.sdk.state.StateSpecs;
  import org.apache.beam.sdk.state.TimeDomain;
  import org.apache.beam.sdk.state.Timer;
  import org.apache.beam.sdk.state.TimerSpec;
  import org.apache.beam.sdk.state.TimerSpecs;
  import org.apache.beam.sdk.state.ValueState;
  import org.apache.beam.sdk.state.WatermarkHoldState;
  import org.apache.beam.sdk.testing.PCollectionViewTesting;
  import org.apache.beam.sdk.transforms.Create;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
  import org.apache.beam.sdk.transforms.View;
  import org.apache.beam.sdk.transforms.join.RawUnionValue;
  import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
  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.PaneInfo;
  import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
  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.PCollectionView;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.sdk.values.WindowedValue;
  import org.apache.beam.sdk.values.WindowedValues;
  import org.apache.beam.sdk.values.WindowedValues.FullWindowedValueCoder;
  import org.apache.beam.sdk.values.WindowingStrategy;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Function;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.FluentIterable;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
  import org.apache.flink.api.common.typeinfo.TypeInformation;
  import org.apache.flink.api.java.functions.KeySelector;
  import org.apache.flink.api.java.typeutils.ValueTypeInfo;
  import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
  import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
  import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
  import org.apache.flink.streaming.util.KeyedTwoInputStreamOperatorTestHarness;
  import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
  import org.apache.flink.streaming.util.TwoInputStreamOperatorTestHarness;
  import org.apache.flink.util.OutputTag;
  import org.checkerframework.checker.nullness.qual.Nullable;
  import org.joda.time.Duration;
  import org.joda.time.Instant;
  import org.junit.Before;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.junit.runners.JUnit4;
  import org.mockito.Mockito;
  import org.powermock.reflect.Whitebox;
  
  /** Tests for {@link DoFnOperator}. */
  @RunWith(JUnit4.class)
  @SuppressWarnings({
    "keyfor",
    // TODO(https://github.com/apache/beam/issues/21230): Remove when new version of
    // errorprone is released (2.11.0)
    "unused"
  })
  public class DoFnOperatorTest {
  
    // views and windows for testing side inputs
    private static final long WINDOW_MSECS_1 = 100;
    private static final long WINDOW_MSECS_2 = 500;
    private PCollectionView<Iterable<String>> view1;
    private PCollectionView<Iterable<String>> view2;
  
    private int numStartBundleCalled = 0;
  
    @Before
    public void setUp() {
      PCollection<String> pc = Pipeline.create().apply(Create.of("1"));
      view1 =
          pc.apply(Window.into(FixedWindows.of(Duration.millis(WINDOW_MSECS_1))))
              .apply(View.asIterable());
      view2 =
          pc.apply(Window.into(FixedWindows.of(Duration.millis(WINDOW_MSECS_2))))
              .apply(View.asIterable());
    }
  
    @Test
    public void testSingleOutput() throws Exception {
  
      Coder<WindowedValue<String>> coder = WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      DoFnOperator<String, String, String> doFnOperator =
          new DoFnOperator<>(
              new IdentityDoFn<>(),
              "stepName",
              coder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag, coder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              FlinkPipelineOptions.defaults(),
              null,
              null,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("Hello")));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(WindowedValues.valueInGlobalWindow("Hello")));
  
      testHarness.close();
    }
  
    @Test
    @SuppressWarnings("unchecked")
    public void testMultiOutputOutput() throws Exception {
  
      WindowedValues.ValueOnlyWindowedValueCoder<String> coder =
          WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
  
      TupleTag<String> mainOutput = new TupleTag<>("main-output");
      TupleTag<String> additionalOutput1 = new TupleTag<>("output-1");
      TupleTag<String> additionalOutput2 = new TupleTag<>("output-2");
      ImmutableMap<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags =
          ImmutableMap.<TupleTag<?>, OutputTag<WindowedValue<?>>>builder()
              .put(additionalOutput1, new OutputTag<WindowedValue<?>>(additionalOutput1.getId()) {})
              .put(additionalOutput2, new OutputTag<WindowedValue<?>>(additionalOutput2.getId()) {})
              .build();
      @SuppressWarnings("rawtypes")
      ImmutableMap<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders =
          ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder()
              .put(mainOutput, (Coder) coder)
              .put(additionalOutput1, coder)
              .put(additionalOutput2, coder)
              .build();
      ImmutableMap<TupleTag<?>, Integer> tagsToIds =
          ImmutableMap.<TupleTag<?>, Integer>builder()
              .put(mainOutput, 0)
              .put(additionalOutput1, 1)
              .put(additionalOutput2, 2)
              .build();
  
      DoFnOperator<String, String, String> doFnOperator =
          new DoFnOperator<>(
              new MultiOutputDoFn(additionalOutput1, additionalOutput2),
              "stepName",
              coder,
              Collections.emptyMap(),
              mainOutput,
              ImmutableList.of(additionalOutput1, additionalOutput2),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  mainOutput,
                  tagsToOutputTags,
                  tagsToCoders,
                  tagsToIds,
                  new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              FlinkPipelineOptions.defaults(),
              null,
              null,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("one")));
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("two")));
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("hello")));
  
      assertThat(
          this.stripStreamRecord(testHarness.getOutput()),
          contains(WindowedValues.valueInGlobalWindow("got: hello")));
  
      assertThat(
          this.stripStreamRecord(testHarness.getSideOutput(tagsToOutputTags.get(additionalOutput1))),
          contains(
              WindowedValues.valueInGlobalWindow("extra: one"),
              WindowedValues.valueInGlobalWindow("got: hello")));
  
      assertThat(
          this.stripStreamRecord(testHarness.getSideOutput(tagsToOutputTags.get(additionalOutput2))),
          contains(
              WindowedValues.valueInGlobalWindow("extra: two"),
              WindowedValues.valueInGlobalWindow("got: hello")));
  
      testHarness.close();
    }
  
    /**
     * This test specifically verifies that we correctly map Flink watermarks to Beam watermarks. In
     * Beam, a watermark {@code T} guarantees there will not be elements with a timestamp {@code < T}
     * in the future. In Flink, a watermark {@code T} guarantees there will not be elements with a
     * timestamp {@code <= T} in the future. We have to make sure to take this into account when
     * firing timers.
     *
     * <p>This does not test the timer API in general or processing-time timers because there are
     * generic tests for this in {@code ParDoTest}.
     */
    @Test
    public void testWatermarkContract() throws Exception {
  
      final Instant timerTimestamp = new Instant(1000);
      final Instant timerOutputTimestamp = timerTimestamp.minus(Duration.millis(1));
      final String eventTimeMessage = "Event timer fired: ";
      final String processingTimeMessage = "Processing timer fired";
  
      WindowingStrategy<Object, IntervalWindow> windowingStrategy =
          WindowingStrategy.of(FixedWindows.of(Duration.millis(10_000)));
  
      final String eventTimerId = "eventTimer";
      final String eventTimerId2 = "eventTimer2";
      final String processingTimerId = "processingTimer";
      DoFn<Integer, String> fn =
          new DoFn<Integer, String>() {
  
            @TimerId(eventTimerId)
            private final TimerSpec eventTimer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
  
            @TimerId(eventTimerId2)
            private final TimerSpec eventTimer2 = TimerSpecs.timer(TimeDomain.EVENT_TIME);
  
            @TimerId(processingTimerId)
            private final TimerSpec processingTimer = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
  
            @ProcessElement
            public void processElement(
                ProcessContext context,
                @TimerId(eventTimerId) Timer eventTimer,
                @TimerId(eventTimerId2) Timer eventTimerWithOutputTimestamp,
                @TimerId(processingTimerId) Timer processingTimer) {
              eventTimer.set(timerTimestamp);
              eventTimerWithOutputTimestamp
                  .withOutputTimestamp(timerOutputTimestamp)
                  .set(timerTimestamp);
              processingTimer
                  .withOutputTimestamp(new Instant(10))
                  .offset(Duration.millis(timerTimestamp.getMillis()))
                  .setRelative();
            }
  
            @OnTimer(eventTimerId)
            public void onEventTime(OnTimerContext context) {
              assertEquals(
                  "Timer timestamp must match set timestamp.", timerTimestamp, context.timestamp());
              context.outputWithTimestamp(eventTimeMessage + eventTimerId, context.timestamp());
            }
  
            @OnTimer(eventTimerId2)
            public void onEventTime2(OnTimerContext context) {
              assertEquals(
                  "Timer timestamp must match set timestamp.",
                  timerTimestamp,
                  context.fireTimestamp());
              context.output(eventTimeMessage + eventTimerId2);
            }
  
            @OnTimer(processingTimerId)
            public void onProcessingTime(OnTimerContext context) {
              assertTrue(
                  // Timestamps in processing timer context are defined to be the output watermark
                  "Timer timestamp must be at most current input watermark",
                  !timerTimestamp.plus(Duration.millis(1)).isBefore(context.timestamp()));
              context.outputWithTimestamp(processingTimeMessage, context.timestamp());
            }
          };
  
      VarIntCoder keyCoder = VarIntCoder.of();
      WindowedValues.FullWindowedValueCoder<Integer> inputCoder =
          WindowedValues.getFullCoder(keyCoder, windowingStrategy.getWindowFn().windowCoder());
  
      WindowedValues.FullWindowedValueCoder<String> outputCoder =
          WindowedValues.getFullCoder(
              StringUtf8Coder.of(), windowingStrategy.getWindowFn().windowCoder());
  
      KeySelector<WindowedValue<Integer>, FlinkKey> keySelector =
          e -> FlinkKey.of(e.getValue(), keyCoder);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      DoFnOperator<Integer, Integer, String> doFnOperator =
          new DoFnOperator<>(
              fn,
              "stepName",
              inputCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag,
                  outputCoder,
                  new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              windowingStrategy,
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              FlinkPipelineOptions.defaults(),
              keyCoder, /* key coder */
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness =
          new KeyedOneInputStreamOperatorTestHarness<>(
              doFnOperator, keySelector, ValueTypeInfo.of(FlinkKey.class));
  
      testHarness.setup(
          new CoderTypeSerializer<>(
              outputCoder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())));
  
      testHarness.open();
  
      testHarness.processWatermark(499);
      testHarness.setProcessingTime(0);
  
      IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.millis(10_000));
  
      // this should register the two timers above
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());
  
      // this does not yet fire the timers (in vanilla Flink it would)
      testHarness.processWatermark(timerTimestamp.getMillis());
      testHarness.setProcessingTime(timerTimestamp.getMillis());
  
      assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());
      assertThat(doFnOperator.keyedStateInternals.minWatermarkHoldMs(), is(10L));
  
      // this must fire the processing timer
      testHarness.setProcessingTime(timerTimestamp.getMillis() + 1);
  
      assertThat(
          doFnOperator.keyedStateInternals.minWatermarkHoldMs(),
          is(timerOutputTimestamp.getMillis()));
  
      // this must fire the event timers
      testHarness.processWatermark(timerTimestamp.getMillis() + 1);
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          containsInAnyOrder(
              WindowedValues.of(
                  eventTimeMessage + eventTimerId, timerTimestamp, window1, PaneInfo.NO_FIRING),
              WindowedValues.of(
                  eventTimeMessage + eventTimerId2,
                  timerTimestamp.minus(Duration.millis(1)),
                  window1,
                  PaneInfo.NO_FIRING),
              WindowedValues.of(
                  processingTimeMessage, new Instant(10), window1, PaneInfo.NO_FIRING)));
  
      testHarness.close();
    }
  
    @Test
    public void testWatermarkUpdateAfterWatermarkHoldRelease() throws Exception {
  
      Coder<WindowedValue<KV<String, String>>> coder =
          WindowedValues.getValueOnlyCoder(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()));
  
      TupleTag<KV<String, String>> outputTag = new TupleTag<>("main-output");
      List<Long> emittedWatermarkHolds = new ArrayList<>();
  
      KeySelector<WindowedValue<KV<String, String>>, FlinkKey> keySelector =
          e -> FlinkKey.of(e.getValue().getKey(), StringUtf8Coder.of());
  
      DoFnOperator<KV<String, String>, KV<String, String>, KV<String, String>> doFnOperator =
          new DoFnOperator<KV<String, String>, KV<String, String>, KV<String, String>>(
              new IdentityDoFn<>(),
              "stepName",
              coder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag, coder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              FlinkPipelineOptions.defaults(),
              StringUtf8Coder.of(),
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap()) {
  
            @Override
            protected DoFnRunner<KV<String, String>, KV<String, String>> createWrappingDoFnRunner(
                DoFnRunner<KV<String, String>, KV<String, String>> wrappedRunner,
                StepContext stepContext) {
  
              StateNamespace namespace =
                  StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE);
              StateTag<WatermarkHoldState> holdTag =
                  StateTags.watermarkStateInternal("hold", TimestampCombiner.LATEST);
              WatermarkHoldState holdState = stepContext.stateInternals().state(namespace, holdTag);
              TimerInternals timerInternals = stepContext.timerInternals();
  
              return new DoFnRunner<KV<String, String>, KV<String, String>>() {
  
                @Override
                public void startBundle() {
                  wrappedRunner.startBundle();
                }
  
                @Override
                public void processElement(WindowedValue<KV<String, String>> elem) {
                  wrappedRunner.processElement(elem);
                  holdState.add(elem.getTimestamp());
                  timerInternals.setTimer(
                      namespace,
                      "timer",
                      "family",
                      elem.getTimestamp().plus(Duration.millis(1)),
                      elem.getTimestamp().plus(Duration.millis(1)),
                      TimeDomain.EVENT_TIME);
                  timerInternals.setTimer(
                      namespace,
                      "cleanup",
                      "",
                      GlobalWindow.INSTANCE.maxTimestamp(),
                      GlobalWindow.INSTANCE.maxTimestamp(),
                      TimeDomain.EVENT_TIME);
                }
  
                @Override
                public <KeyT> void onTimer(
                    String timerId,
                    String timerFamilyId,
                    KeyT key,
                    BoundedWindow window,
                    Instant timestamp,
                    Instant outputTimestamp,
                    TimeDomain timeDomain) {
  
                  if ("cleanup".equals(timerId)) {
                    holdState.clear();
                  } else {
                    holdState.add(outputTimestamp);
                  }
                }
  
                @Override
                public void finishBundle() {
                  wrappedRunner.finishBundle();
                }
  
                @Override
                public <KeyT> void onWindowExpiration(
                    BoundedWindow window, Instant timestamp, KeyT key) {
                  wrappedRunner.onWindowExpiration(window, timestamp, key);
                }
  
                @Override
                public DoFn<KV<String, String>, KV<String, String>> getFn() {
                  return doFn;
                }
              };
            }
  
            @Override
            void emitWatermarkIfHoldChanged(long currentWatermarkHold) {
              emittedWatermarkHolds.add(keyedStateInternals.minWatermarkHoldMs());
            }
          };
  
      OneInputStreamOperatorTestHarness<
              WindowedValue<KV<String, String>>, WindowedValue<KV<String, String>>>
          testHarness =
              new KeyedOneInputStreamOperatorTestHarness<>(
                  doFnOperator, keySelector, ValueTypeInfo.of(FlinkKey.class));
  
      testHarness.setup();
  
      Instant now = Instant.now();
  
      testHarness.open();
  
      // process first element, set hold to `now', setup timer for `now + 1'
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.timestampedValueInGlobalWindow(KV.of("Key", "Hello"), now)));
  
      assertThat(emittedWatermarkHolds, is(equalTo(Collections.singletonList(now.getMillis()))));
  
      // fire timer, change hold to `now + 2'
      testHarness.processWatermark(now.getMillis() + 2);
  
      assertThat(
          emittedWatermarkHolds, is(equalTo(Arrays.asList(now.getMillis(), now.getMillis() + 1))));
  
      // process second element, verify we emitted changed hold
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.timestampedValueInGlobalWindow(
                  KV.of("Key", "Hello"), now.plus(Duration.millis(2)))));
  
      assertThat(
          emittedWatermarkHolds,
          is(equalTo(Arrays.asList(now.getMillis(), now.getMillis() + 1, now.getMillis() + 2))));
  
      testHarness.processWatermark(
          GlobalWindow.INSTANCE.maxTimestamp().plus(Duration.millis(1)).getMillis());
      testHarness.processWatermark(BoundedWindow.TIMESTAMP_MAX_VALUE.getMillis());
  
      testHarness.close();
    }
  
    @Test
    public void testLateDroppingForStatefulFn() throws Exception {
  
      WindowingStrategy<Object, IntervalWindow> windowingStrategy =
          WindowingStrategy.of(FixedWindows.of(Duration.millis(10)));
  
      DoFn<Integer, String> fn =
          new DoFn<Integer, String>() {
  
            @StateId("state")
            private final StateSpec<ValueState<String>> stateSpec =
                StateSpecs.value(StringUtf8Coder.of());
  
            @ProcessElement
            public void processElement(ProcessContext context) {
              context.output(context.element().toString());
            }
          };
  
      VarIntCoder keyCoder = VarIntCoder.of();
      Coder<WindowedValue<Integer>> inputCoder =
          WindowedValues.getFullCoder(keyCoder, windowingStrategy.getWindowFn().windowCoder());
      Coder<WindowedValue<String>> outputCoder =
          WindowedValues.getFullCoder(
              StringUtf8Coder.of(), windowingStrategy.getWindowFn().windowCoder());
  
      KeySelector<WindowedValue<Integer>, FlinkKey> keySelector =
          e -> FlinkKey.of(e.getValue(), keyCoder);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      DoFnOperator<Integer, Integer, String> doFnOperator =
          new DoFnOperator<>(
              fn,
              "stepName",
              inputCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag,
                  outputCoder,
                  new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              windowingStrategy,
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              FlinkPipelineOptions.defaults(),
              keyCoder, /* key coder */
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness =
          new KeyedOneInputStreamOperatorTestHarness<>(
              doFnOperator, keySelector, ValueTypeInfo.of(FlinkKey.class));
  
      testHarness.open();
  
      testHarness.processWatermark(0);
  
      IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.millis(10));
  
      // this should not be late
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(WindowedValues.of("13", new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      testHarness.getOutput().clear();
  
      testHarness.processWatermark(9);
  
      // this should still not be considered late
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(WindowedValues.of("17", new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      testHarness.getOutput().clear();
  
      testHarness.processWatermark(10);
  
      // this should now be considered late
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.of(17, new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());
  
      testHarness.close();
    }
  
    @Test
    public void testStateGCForStatefulFn() throws Exception {
  
      WindowingStrategy<Object, IntervalWindow> windowingStrategy =
          WindowingStrategy.of(FixedWindows.of(Duration.millis(10)))
              .withAllowedLateness(Duration.ZERO);
      final int offset = 5000;
      final int timerOutput = 4093;
  
      KeyedOneInputStreamOperatorTestHarness<
              FlinkKey, WindowedValue<KV<String, Integer>>, WindowedValue<KV<String, Integer>>>
          testHarness =
              getHarness(
                  windowingStrategy,
                  offset,
                  (window) -> Objects.requireNonNull(window).maxTimestamp(),
                  timerOutput);
  
      testHarness.open();
  
      testHarness.processWatermark(0);
  
      assertEquals(0, testHarness.numKeyedStateEntries());
  
      IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.millis(10));
  
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.of(KV.of("key1", 5), new Instant(1), window1, PaneInfo.NO_FIRING)));
  
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.of(KV.of("key2", 7), new Instant(3), window1, PaneInfo.NO_FIRING)));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.of(
                  KV.of("key1", 5 + offset), new Instant(1), window1, PaneInfo.NO_FIRING),
              WindowedValues.of(
                  KV.of("key2", 7 + offset), new Instant(3), window1, PaneInfo.NO_FIRING)));
  
      // 2 entries for the elements and 2 for the pending timers
      assertEquals(4, testHarness.numKeyedStateEntries());
  
      testHarness.getOutput().clear();
  
      // this should trigger both the window.maxTimestamp() timer and the GC timer
      // this tests that the GC timer fires after the user timer
      // we have to add 1 here because Flink timers fire when watermark >= timestamp while Beam
      // timers fire when watermark > timestamp
      testHarness.processWatermark(
          window1
                  .maxTimestamp()
                  .plus(windowingStrategy.getAllowedLateness())
                  .plus(Duration.millis(StatefulDoFnRunner.TimeInternalsCleanupTimer.GC_DELAY_MS))
                  .getMillis()
              + 1);
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.of(
                  KV.of("key1", timerOutput), new Instant(9), window1, PaneInfo.NO_FIRING),
              WindowedValues.of(
                  KV.of("key2", timerOutput), new Instant(9), window1, PaneInfo.NO_FIRING)));
  
      testHarness.close();
    }
  
    @Test
    public void testGCForGlobalWindow() throws Exception {
      WindowingStrategy<Object, GlobalWindow> windowingStrategy = WindowingStrategy.globalDefault();
  
      KeyedOneInputStreamOperatorTestHarness<
              FlinkKey, WindowedValue<KV<String, Integer>>, WindowedValue<KV<String, Integer>>>
          testHarness = getHarness(windowingStrategy, 5000, (window) -> new Instant(50), 4092);
  
      testHarness.open();
  
      testHarness.processWatermark(0);
  
      // ensure the state was garbage collected and the pending timers have been removed
      assertEquals(0, testHarness.numKeyedStateEntries());
  
      // Check global window cleanup via final watermark, _without_ cleanup timers
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.of(
                  KV.of("key1", 5), new Instant(23), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)));
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.of(
                  KV.of("key2", 6), new Instant(42), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)));
  
      // timers set by the transform
      assertThat(testHarness.numEventTimeTimers(), is(2));
      // state has been written which needs to be cleaned up (includes timers)
      assertThat(testHarness.numKeyedStateEntries(), is(4));
  
      // Fire timers set
      testHarness.processWatermark(51);
      assertThat(testHarness.numEventTimeTimers(), is(0));
      assertThat(testHarness.numKeyedStateEntries(), is(2));
  
      // enforce closing of bundle
      testHarness.setProcessingTime(
          testHarness.getProcessingTime()
              + 2 * FlinkPipelineOptions.defaults().getMaxBundleTimeMills());
  
      // Should not trigger garbage collection yet
      testHarness.processWatermark(
          GlobalWindow.INSTANCE.maxTimestamp().plus(Duration.millis(1)).getMillis());
      assertThat(testHarness.numEventTimeTimers(), is(0));
      assertThat(testHarness.numKeyedStateEntries(), is(2));
  
      // Cleanup due to end of global window
      testHarness.processWatermark(
          GlobalWindow.INSTANCE.maxTimestamp().plus(Duration.millis(2)).getMillis());
      assertThat(testHarness.numEventTimeTimers(), is(0));
      assertThat(testHarness.numKeyedStateEntries(), is(0));
  
      // Any new state will also be cleaned up on close
      testHarness.processElement(
          new StreamRecord<>(
              WindowedValues.of(
                  KV.of("key2", 6), new Instant(42), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)));
  
      // Close sends Flink's max watermark and will cleanup again
      testHarness.close();
      assertThat(testHarness.numEventTimeTimers(), is(0));
      assertThat(testHarness.numKeyedStateEntries(), is(0));
    }
  
    private static KeyedOneInputStreamOperatorTestHarness<
            FlinkKey, WindowedValue<KV<String, Integer>>, WindowedValue<KV<String, Integer>>>
        getHarness(
            WindowingStrategy<?, ?> windowingStrategy,
            int elementOffset,
            Function<BoundedWindow, Instant> timerTimestamp,
            int timerOutput)
            throws Exception {
      final String timerId = "boo";
      final String stateId = "dazzle";
  
      DoFn<KV<String, Integer>, KV<String, Integer>> fn =
          new DoFn<KV<String, Integer>, KV<String, Integer>>() {
  
            @TimerId(timerId)
            private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
  
            @StateId(stateId)
            private final StateSpec<ValueState<String>> stateSpec =
                StateSpecs.value(StringUtf8Coder.of());
  
            @ProcessElement
            public void processElement(
                ProcessContext context,
                @TimerId(timerId) Timer timer,
                @StateId(stateId) ValueState<String> state,
                BoundedWindow window) {
              timer.set(Objects.requireNonNull(timerTimestamp.apply(window)));
              state.write(context.element().getKey());
              context.output(
                  KV.of(context.element().getKey(), context.element().getValue() + elementOffset));
            }
  
            @OnTimer(timerId)
            public void onTimer(OnTimerContext context, @StateId(stateId) ValueState<String> state) {
              context.output(KV.of(state.read(), timerOutput));
            }
          };
  
      WindowedValues.FullWindowedValueCoder<KV<String, Integer>> coder =
          WindowedValues.getFullCoder(
              KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()),
              windowingStrategy.getWindowFn().windowCoder());
  
      TupleTag<KV<String, Integer>> outputTag = new TupleTag<>("main-output");
  
      KeySelector<WindowedValue<KV<String, Integer>>, FlinkKey> keySelector =
          e -> FlinkKey.of(e.getValue().getKey(), StringUtf8Coder.of());
  
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setStreaming(true);
  
      DoFnOperator<KV<String, Integer>, KV<String, Integer>, KV<String, Integer>> doFnOperator =
          new DoFnOperator<>(
              fn,
              "stepName",
              coder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag, coder, new SerializablePipelineOptions(options)),
              windowingStrategy,
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              options,
              StringUtf8Coder.of(), /* key coder */
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      return new KeyedOneInputStreamOperatorTestHarness<>(
          doFnOperator, keySelector, ValueTypeInfo.of(FlinkKey.class));
    }
  
    @Test
    public void testNormalParDoSideInputs() throws Exception {
      testSideInputs(false);
    }
  
    @Test
    public void testKeyedParDoSideInputs() throws Exception {
      testSideInputs(true);
    }
  
    void testSideInputs(boolean keyed) throws Exception {
  
      Coder<WindowedValue<String>> coder = WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      ImmutableMap<Integer, PCollectionView<?>> sideInputMapping =
          ImmutableMap.<Integer, PCollectionView<?>>builder().put(1, view1).put(2, view2).build();
  
      Coder<String> keyCoder = StringUtf8Coder.of();
      KeySelector<WindowedValue<String>, FlinkKey> keySelector = null;
      if (keyed) {
        keySelector = value -> FlinkKey.of(value.getValue(), keyCoder);
      }
  
      DoFnOperator<String, String, String> doFnOperator =
          new DoFnOperator<>(
              new IdentityDoFn<>(),
              "stepName",
              coder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag, coder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              WindowingStrategy.of(FixedWindows.of(Duration.millis(100))),
              sideInputMapping, /* side-input mapping */
              ImmutableList.of(view1, view2), /* side inputs */
              FlinkPipelineOptions.defaults(),
              keyed ? keyCoder : null,
              keyed ? keySelector : null,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      TwoInputStreamOperatorTestHarness<WindowedValue<String>, RawUnionValue, WindowedValue<String>>
          testHarness = new TwoInputStreamOperatorTestHarness<>(doFnOperator);
  
      if (keyed) {
        // we use a dummy key for the second input since it is considered to be broadcast
        testHarness =
            new KeyedTwoInputStreamOperatorTestHarness<>(
                doFnOperator, keySelector, null, ValueTypeInfo.of(FlinkKey.class));
      }
  
      testHarness.open();
  
      IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(100));
      IntervalWindow secondWindow = new IntervalWindow(new Instant(0), new Instant(500));
  
      // test the keep of sideInputs events
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  1,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view1.getPipeline().getOptions(), View.asIterable(), "hello", "ciao"),
                      new Instant(0),
                      firstWindow))));
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  2,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view2.getPipeline().getOptions(), View.asIterable(), "foo", "bar"),
                      new Instant(0),
                      secondWindow))));
  
      // push in a regular elements
      WindowedValue<String> helloElement = valueInWindow("Hello", new Instant(0), firstWindow);
      WindowedValue<String> worldElement = valueInWindow("World", new Instant(1000), firstWindow);
      testHarness.processElement1(new StreamRecord<>(helloElement));
      testHarness.processElement1(new StreamRecord<>(worldElement));
  
      // test the keep of pushed-back events
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  1,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view1.getPipeline().getOptions(), View.asIterable(), "hello", "ciao"),
                      new Instant(1000),
                      firstWindow))));
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  2,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view2.getPipeline().getOptions(), View.asIterable(), "foo", "bar"),
                      new Instant(1000),
                      secondWindow))));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(helloElement, worldElement));
  
      testHarness.close();
    }
  
    @Test
    public void testStateRestore() throws Exception {
      DoFn<KV<String, Long>, KV<String, Long>> filterElementsEqualToCountFn =
          new DoFn<KV<String, Long>, KV<String, Long>>() {
  
            @StateId("counter")
            private final StateSpec<ValueState<Long>> counterSpec =
                StateSpecs.value(VarLongCoder.of());
  
            @ProcessElement
            public void processElement(
                ProcessContext context, @StateId("counter") ValueState<Long> count) {
              long currentCount = Optional.ofNullable(count.read()).orElse(0L);
              currentCount = currentCount + 1;
              count.write(currentCount);
  
              KV<String, Long> currentElement = context.element();
              if (currentCount == currentElement.getValue()) {
                context.output(currentElement);
              }
            }
          };
  
      WindowingStrategy<Object, GlobalWindow> windowingStrategy = WindowingStrategy.globalDefault();
  
      TupleTag<KV<String, Long>> outputTag = new TupleTag<>("main-output");
  
      StringUtf8Coder keyCoder = StringUtf8Coder.of();
      KvToFlinkKeyKeySelector<String, Long> keySelector = new KvToFlinkKeyKeySelector<>(keyCoder);
  
      KvCoder<String, Long> coder = KvCoder.of(keyCoder, VarLongCoder.of());
  
      FullWindowedValueCoder<KV<String, Long>> kvCoder =
          WindowedValues.getFullCoder(coder, windowingStrategy.getWindowFn().windowCoder());
  
      TypeInformation<FlinkKey> keyCoderInfo = ValueTypeInfo.of(FlinkKey.class);
  
      OneInputStreamOperatorTestHarness<
              WindowedValue<KV<String, Long>>, WindowedValue<KV<String, Long>>>
          testHarness =
              createTestHarness(
                  windowingStrategy,
                  filterElementsEqualToCountFn,
                  kvCoder,
                  kvCoder,
                  keyCoder,
                  outputTag,
                  keyCoderInfo,
                  keySelector);
      testHarness.open();
  
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("a", 100L))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("a", 100L))));
  
      OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
      testHarness.close();
  
      testHarness =
          createTestHarness(
              windowingStrategy,
              filterElementsEqualToCountFn,
              kvCoder,
              kvCoder,
              keyCoder,
              outputTag,
              keyCoderInfo,
              keySelector);
      testHarness.initializeState(snapshot);
      testHarness.open();
  
      // after restore: counter = 2
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("a", 100L))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("a", 4L))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("a", 5L))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("a", 100L))));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow(KV.of("a", 4L)),
              WindowedValues.valueInGlobalWindow(KV.of("a", 5L))));
  
      testHarness.close();
    }
  
    @Test
    public void nonKeyedParDoSideInputCheckpointing() throws Exception {
      sideInputCheckpointing(
          () -> {
            Coder<WindowedValue<String>> coder =
                WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder());
            TupleTag<String> outputTag = new TupleTag<>("main-output");
  
            ImmutableMap<Integer, PCollectionView<?>> sideInputMapping =
                ImmutableMap.<Integer, PCollectionView<?>>builder()
                    .put(1, view1)
                    .put(2, view2)
                    .build();
  
            DoFnOperator<String, String, String> doFnOperator =
                new DoFnOperator<>(
                    new IdentityDoFn<>(),
                    "stepName",
                    coder,
                    Collections.emptyMap(),
                    outputTag,
                    Collections.emptyList(),
                    new DoFnOperator.MultiOutputOutputManagerFactory<>(
                        outputTag,
                        coder,
                        new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
                    WindowingStrategy.globalDefault(),
                    sideInputMapping, /* side-input mapping */
                    ImmutableList.of(view1, view2), /* side inputs */
                    FlinkPipelineOptions.defaults(),
                    null,
                    null,
                    DoFnSchemaInformation.create(),
                    Collections.emptyMap());
  
            return new TwoInputStreamOperatorTestHarness<>(doFnOperator);
          });
    }
  
    @Test
    public void keyedParDoSideInputCheckpointing() throws Exception {
      sideInputCheckpointing(
          () -> {
            StringUtf8Coder keyCoder = StringUtf8Coder.of();
            Coder<WindowedValue<String>> coder =
                WindowedValues.getFullCoder(keyCoder, IntervalWindow.getCoder());
            TupleTag<String> outputTag = new TupleTag<>("main-output");
  
            KeySelector<WindowedValue<String>, FlinkKey> keySelector =
                e -> FlinkKey.of(e.getValue(), keyCoder);
  
            ImmutableMap<Integer, PCollectionView<?>> sideInputMapping =
                ImmutableMap.<Integer, PCollectionView<?>>builder()
                    .put(1, view1)
                    .put(2, view2)
                    .build();
  
            DoFnOperator<String, String, String> doFnOperator =
                new DoFnOperator<>(
                    new IdentityDoFn<>(),
                    "stepName",
                    coder,
                    Collections.emptyMap(),
                    outputTag,
                    Collections.emptyList(),
                    new DoFnOperator.MultiOutputOutputManagerFactory<>(
                        outputTag,
                        coder,
                        new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
                    WindowingStrategy.of(FixedWindows.of(Duration.millis(100))),
                    sideInputMapping, /* side-input mapping */
                    ImmutableList.of(view1, view2), /* side inputs */
                    FlinkPipelineOptions.defaults(),
                    keyCoder,
                    keySelector,
                    DoFnSchemaInformation.create(),
                    Collections.emptyMap());
  
            return new KeyedTwoInputStreamOperatorTestHarness<>(
                doFnOperator,
                keySelector,
                // we use a dummy key for the second input since it is considered to be broadcast
                null,
                ValueTypeInfo.of(FlinkKey.class));
          });
    }
  
    void sideInputCheckpointing(
        TestHarnessFactory<
                TwoInputStreamOperatorTestHarness<
                    WindowedValue<String>, RawUnionValue, WindowedValue<String>>>
            harnessFactory)
        throws Exception {
  
      TwoInputStreamOperatorTestHarness<WindowedValue<String>, RawUnionValue, WindowedValue<String>>
          testHarness = harnessFactory.create();
  
      testHarness.open();
  
      IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(100));
      IntervalWindow secondWindow = new IntervalWindow(new Instant(0), new Instant(500));
  
      // push in some side inputs for both windows
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  1,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view1.getPipeline().getOptions(), View.asIterable(), "hello", "ciao"),
                      new Instant(0),
                      firstWindow))));
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  2,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view2.getPipeline().getOptions(), View.asIterable(), "foo", "bar"),
                      new Instant(0),
                      secondWindow))));
  
      // snapshot state, throw away the operator, then restore and verify that we still match
      // main-input elements to the side-inputs that we sent earlier
      OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
  
      testHarness = harnessFactory.create();
  
      testHarness.initializeState(snapshot);
      testHarness.open();
  
      // push in main-input elements
      WindowedValue<String> helloElement = valueInWindow("Hello", new Instant(0), firstWindow);
      WindowedValue<String> worldElement = valueInWindow("World", new Instant(1000), firstWindow);
      testHarness.processElement1(new StreamRecord<>(helloElement));
      testHarness.processElement1(new StreamRecord<>(worldElement));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(helloElement, worldElement));
  
      testHarness.close();
    }
  
    @Test
    public void nonKeyedParDoPushbackDataCheckpointing() throws Exception {
      pushbackDataCheckpointing(
          () -> {
            Coder<WindowedValue<String>> coder =
                WindowedValues.getFullCoder(StringUtf8Coder.of(), IntervalWindow.getCoder());
  
            TupleTag<String> outputTag = new TupleTag<>("main-output");
  
            ImmutableMap<Integer, PCollectionView<?>> sideInputMapping =
                ImmutableMap.<Integer, PCollectionView<?>>builder()
                    .put(1, view1)
                    .put(2, view2)
                    .build();
  
            DoFnOperator<String, String, String> doFnOperator =
                new DoFnOperator<>(
                    new IdentityDoFn<>(),
                    "stepName",
                    coder,
                    Collections.emptyMap(),
                    outputTag,
                    Collections.emptyList(),
                    new DoFnOperator.MultiOutputOutputManagerFactory<>(
                        outputTag,
                        coder,
                        new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
                    WindowingStrategy.of(FixedWindows.of(Duration.millis(100))),
                    sideInputMapping, /* side-input mapping */
                    ImmutableList.of(view1, view2), /* side inputs */
                    FlinkPipelineOptions.defaults(),
                    null,
                    null,
                    DoFnSchemaInformation.create(),
                    Collections.emptyMap());
  
            return new TwoInputStreamOperatorTestHarness<>(doFnOperator);
          });
    }
  
    @Test
    public void keyedParDoPushbackDataCheckpointing() throws Exception {
      pushbackDataCheckpointing(
          () -> {
            StringUtf8Coder keyCoder = StringUtf8Coder.of();
            Coder<WindowedValue<String>> coder =
                WindowedValues.getFullCoder(keyCoder, IntervalWindow.getCoder());
  
            TupleTag<String> outputTag = new TupleTag<>("main-output");
  
            KeySelector<WindowedValue<String>, FlinkKey> keySelector =
                e -> FlinkKey.of(e.getValue(), keyCoder);
  
            ImmutableMap<Integer, PCollectionView<?>> sideInputMapping =
                ImmutableMap.<Integer, PCollectionView<?>>builder()
                    .put(1, view1)
                    .put(2, view2)
                    .build();
  
            DoFnOperator<String, String, String> doFnOperator =
                new DoFnOperator<>(
                    new IdentityDoFn<>(),
                    "stepName",
                    coder,
                    Collections.emptyMap(),
                    outputTag,
                    Collections.emptyList(),
                    new DoFnOperator.MultiOutputOutputManagerFactory<>(
                        outputTag,
                        coder,
                        new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
                    WindowingStrategy.of(FixedWindows.of(Duration.millis(100))),
                    sideInputMapping, /* side-input mapping */
                    ImmutableList.of(view1, view2), /* side inputs */
                    FlinkPipelineOptions.defaults(),
                    keyCoder,
                    keySelector,
                    DoFnSchemaInformation.create(),
                    Collections.emptyMap());
  
            return new KeyedTwoInputStreamOperatorTestHarness<>(
                doFnOperator,
                keySelector,
                // we use a dummy key for the second input since it is considered to be broadcast
                null,
                ValueTypeInfo.of(FlinkKey.class));
          });
    }
  
    void pushbackDataCheckpointing(
        TestHarnessFactory<
                TwoInputStreamOperatorTestHarness<
                    WindowedValue<String>, RawUnionValue, WindowedValue<String>>>
            harnessFactory)
        throws Exception {
  
      TwoInputStreamOperatorTestHarness<WindowedValue<String>, RawUnionValue, WindowedValue<String>>
          testHarness = harnessFactory.create();
  
      testHarness.open();
  
      IntervalWindow firstWindow = new IntervalWindow(new Instant(0), new Instant(100));
      IntervalWindow secondWindow = new IntervalWindow(new Instant(0), new Instant(500));
  
      // push in main-input elements
      WindowedValue<String> helloElement = valueInWindow("Hello", new Instant(0), firstWindow);
      WindowedValue<String> worldElement = valueInWindow("World", new Instant(1000), firstWindow);
      testHarness.processElement1(new StreamRecord<>(helloElement));
      testHarness.processElement1(new StreamRecord<>(worldElement));
  
      // snapshot state, throw away the operator, then restore and verify that we still match
      // main-input elements to the side-inputs that we sent earlier
      OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
  
      testHarness = harnessFactory.create();
  
      testHarness.initializeState(snapshot);
      testHarness.open();
  
      // push in some side inputs for both windows
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  1,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view1.getPipeline().getOptions(), View.asIterable(), "hello", "ciao"),
                      new Instant(0),
                      firstWindow))));
      testHarness.processElement2(
          new StreamRecord<>(
              new RawUnionValue(
                  2,
                  valuesInWindow(
                      PCollectionViewTesting.materializeValuesFor(
                          view2.getPipeline().getOptions(), View.asIterable(), "foo", "bar"),
                      new Instant(0),
                      secondWindow))));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          containsInAnyOrder(helloElement, worldElement));
  
      testHarness.close();
    }
  
    @Test
    public void testTimersRestore() throws Exception {
      final Instant timerTimestamp = new Instant(1000);
      final String outputMessage = "Timer fired";
  
      WindowingStrategy<Object, IntervalWindow> windowingStrategy =
          WindowingStrategy.of(FixedWindows.of(Duration.millis(10_000)));
  
      DoFn<Integer, String> fn =
          new DoFn<Integer, String>() {
            private static final String EVENT_TIMER_ID = "eventTimer";
  
            @TimerId(EVENT_TIMER_ID)
            private final TimerSpec eventTimer = TimerSpecs.timer(TimeDomain.EVENT_TIME);
  
            @ProcessElement
            public void processElement(ProcessContext context, @TimerId(EVENT_TIMER_ID) Timer timer) {
              timer.set(timerTimestamp);
            }
  
            @OnTimer(EVENT_TIMER_ID)
            public void onEventTime(OnTimerContext context) {
              assertEquals(
                  "Timer timestamp must match set timestamp.", timerTimestamp, context.timestamp());
              context.outputWithTimestamp(outputMessage, context.timestamp());
            }
          };
  
      VarIntCoder keyCoder = VarIntCoder.of();
      WindowedValues.FullWindowedValueCoder<Integer> inputCoder =
          WindowedValues.getFullCoder(keyCoder, windowingStrategy.getWindowFn().windowCoder());
  
      WindowedValues.FullWindowedValueCoder<String> outputCoder =
          WindowedValues.getFullCoder(
              StringUtf8Coder.of(), windowingStrategy.getWindowFn().windowCoder());
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      final CoderTypeSerializer<WindowedValue<String>> outputSerializer =
          new CoderTypeSerializer<>(
              outputCoder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults()));
      TypeInformation<FlinkKey> keyCoderInfo = ValueTypeInfo.of(FlinkKey.class);
  
      KeySelector<WindowedValue<Integer>, FlinkKey> keySelector =
          e -> FlinkKey.of(e.getValue(), keyCoder);
  
      OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<String>> testHarness =
          createTestHarness(
              windowingStrategy,
              fn,
              inputCoder,
              outputCoder,
              keyCoder,
              outputTag,
              keyCoderInfo,
              keySelector);
  
      testHarness.setup(outputSerializer);
  
      testHarness.open();
  
      testHarness.processWatermark(0);
  
      IntervalWindow window1 = new IntervalWindow(new Instant(0), Duration.millis(10_000));
  
      // this should register a timer
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.of(13, new Instant(0), window1, PaneInfo.NO_FIRING)));
  
      assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());
  
      // snapshot and restore
      final OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
      testHarness.close();
  
      testHarness =
          createTestHarness(
              windowingStrategy,
              fn,
              inputCoder,
              outputCoder,
              VarIntCoder.of(),
              outputTag,
              keyCoderInfo,
              keySelector);
      testHarness.setup(outputSerializer);
      testHarness.initializeState(snapshot);
      testHarness.open();
  
      // this must fire the timer
      testHarness.processWatermark(timerTimestamp.getMillis() + 1);
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(WindowedValues.of(outputMessage, timerTimestamp, window1, PaneInfo.NO_FIRING)));
  
      testHarness.close();
    }
  
    private <K, InT, OutT>
        OneInputStreamOperatorTestHarness<WindowedValue<InT>, WindowedValue<OutT>> createTestHarness(
            WindowingStrategy<Object, ?> windowingStrategy,
            DoFn<InT, OutT> fn,
            FullWindowedValueCoder<InT> inputCoder,
            FullWindowedValueCoder<OutT> outputCoder,
            Coder<?> keyCoder,
            TupleTag<OutT> outputTag,
            TypeInformation<K> keyCoderInfo,
            KeySelector<WindowedValue<InT>, K> keySelector)
            throws Exception {
      DoFnOperator<InT, InT, OutT> doFnOperator =
          new DoFnOperator<>(
              fn,
              "stepName",
              inputCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              new DoFnOperator.MultiOutputOutputManagerFactory<>(
                  outputTag,
                  outputCoder,
                  new SerializablePipelineOptions(FlinkPipelineOptions.defaults())),
              windowingStrategy,
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              FlinkPipelineOptions.defaults(),
              keyCoder /* key coder */,
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      return new KeyedOneInputStreamOperatorTestHarness<>(doFnOperator, keySelector, keyCoderInfo);
    }
  
    @Test
    public void testBundle() throws Exception {
  
      WindowedValues.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(2L);
      options.setMaxBundleTimeMills(10L);
      options.setStreaming(true);
  
      IdentityDoFn<String> doFn =
          new IdentityDoFn<String>() {
            @FinishBundle
            public void finishBundle(FinishBundleContext context) {
              context.output(
                  "finishBundle", BoundedWindow.TIMESTAMP_MIN_VALUE, GlobalWindow.INSTANCE);
            }
          };
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      DoFnOperator<String, String, String> doFnOperator =
          new DoFnOperator<>(
              doFn,
              "stepName",
              windowedValueCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              outputManagerFactory,
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              options,
              null,
              null,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("a")));
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("b")));
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("c")));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("a"),
              WindowedValues.valueInGlobalWindow("b"),
              WindowedValues.valueInGlobalWindow("finishBundle"),
              WindowedValues.valueInGlobalWindow("c")));
  
      // draw a snapshot
      OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
  
      // Finish bundle element will be buffered as part of finishing a bundle in snapshot()
      PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler =
          doFnOperator.outputManager.pushedBackElementsHandler;
      assertThat(pushedBackElementsHandler, instanceOf(NonKeyedPushedBackElementsHandler.class));
      List<KV<Integer, WindowedValue<?>>> bufferedElements =
          pushedBackElementsHandler.getElements().collect(Collectors.toList());
      assertThat(
          bufferedElements, contains(KV.of(0, WindowedValues.valueInGlobalWindow("finishBundle"))));
  
      testHarness.close();
  
      DoFnOperator<String, String, String> newDoFnOperator =
          new DoFnOperator<>(
              doFn,
              "stepName",
              windowedValueCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              outputManagerFactory,
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              options,
              null,
              null,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> newHarness =
          new OneInputStreamOperatorTestHarness<>(newDoFnOperator);
  
      // restore snapshot
      newHarness.initializeState(snapshot);
  
      newHarness.open();
  
      // startBundle will output the buffered elements.
      newHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("d")));
  
      // check finishBundle by timeout
      newHarness.setProcessingTime(10);
  
      assertThat(
          stripStreamRecordFromWindowedValue(newHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("finishBundle"),
              WindowedValues.valueInGlobalWindow("d"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
  
      // No bundle will be created when sending the MAX watermark
      // (unless pushed back items are emitted)
      newHarness.close();
  
      assertThat(
          stripStreamRecordFromWindowedValue(newHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("finishBundle"),
              WindowedValues.valueInGlobalWindow("d"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
  
      // close() will also call dispose(), but call again to verify no new bundle
      // is created afterwards
      newDoFnOperator.cleanUp();
  
      assertThat(
          stripStreamRecordFromWindowedValue(newHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("finishBundle"),
              WindowedValues.valueInGlobalWindow("d"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
    }
  
    @Test
    public void testBundleKeyed() throws Exception {
  
      StringUtf8Coder keyCoder = StringUtf8Coder.of();
      KvToFlinkKeyKeySelector<String, String> keySelector = new KvToFlinkKeyKeySelector<>(keyCoder);
      KvCoder<String, String> kvCoder = KvCoder.of(keyCoder, StringUtf8Coder.of());
      WindowedValues.ValueOnlyWindowedValueCoder<KV<String, String>> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(kvCoder);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(2L);
      options.setMaxBundleTimeMills(10L);
      options.setStreaming(true);
  
      DoFn<KV<String, String>, String> doFn =
          new DoFn<KV<String, String>, String>() {
            @ProcessElement
            public void processElement(ProcessContext ctx) {
              // Change output type of element to test that we do not depend on the input keying
              ctx.output(ctx.element().getValue());
            }
  
            @FinishBundle
            public void finishBundle(FinishBundleContext context) {
              context.output(
                  "finishBundle", BoundedWindow.TIMESTAMP_MIN_VALUE, GlobalWindow.INSTANCE);
            }
          };
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(kvCoder.getValueCoder(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      DoFnOperator<KV<String, String>, KV<String, String>, String> doFnOperator =
          new DoFnOperator<>(
              doFn,
              "stepName",
              windowedValueCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              outputManagerFactory,
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              options,
              keyCoder,
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<KV<String, String>>, WindowedValue<String>>
          testHarness =
              new KeyedOneInputStreamOperatorTestHarness<>(
                  doFnOperator, keySelector, keySelector.getProducedType());
  
      testHarness.open();
  
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key", "a"))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key", "b"))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key", "c"))));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("a"),
              WindowedValues.valueInGlobalWindow("b"),
              WindowedValues.valueInGlobalWindow("finishBundle"),
              WindowedValues.valueInGlobalWindow("c")));
  
      // Take a snapshot
      OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
  
      // Finish bundle element will be buffered as part of finishing a bundle in snapshot()
      PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler =
          doFnOperator.outputManager.pushedBackElementsHandler;
      assertThat(pushedBackElementsHandler, instanceOf(NonKeyedPushedBackElementsHandler.class));
      List<KV<Integer, WindowedValue<?>>> bufferedElements =
          pushedBackElementsHandler.getElements().collect(Collectors.toList());
      assertThat(
          bufferedElements, contains(KV.of(0, WindowedValues.valueInGlobalWindow("finishBundle"))));
  
      testHarness.close();
  
      doFnOperator =
          new DoFnOperator<>(
              doFn,
              "stepName",
              windowedValueCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              outputManagerFactory,
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              options,
              keyCoder,
              keySelector,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      testHarness =
          new KeyedOneInputStreamOperatorTestHarness<>(
              doFnOperator, keySelector, keySelector.getProducedType());
  
      // Restore snapshot
      testHarness.initializeState(snapshot);
  
      testHarness.open();
  
      // startBundle will output the buffered elements.
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key", "d"))));
  
      // check finishBundle by timeout
      testHarness.setProcessingTime(10);
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              // The first finishBundle is restored from the checkpoint
              WindowedValues.valueInGlobalWindow("finishBundle"),
              WindowedValues.valueInGlobalWindow("d"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
  
      testHarness.close();
    }
  
    @Test
    public void testCheckpointBufferingWithMultipleBundles() throws Exception {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(10L);
      options.setCheckpointingInterval(1L);
      options.setStreaming(true);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      StringUtf8Coder coder = StringUtf8Coder.of();
      WindowedValues.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(coder);
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      Supplier<DoFnOperator<String, String, String>> doFnOperatorSupplier =
          () ->
              new DoFnOperator<>(
                  new IdentityDoFn<>(),
                  "stepName",
                  windowedValueCoder,
                  Collections.emptyMap(),
                  outputTag,
                  Collections.emptyList(),
                  outputManagerFactory,
                  WindowingStrategy.globalDefault(),
                  new HashMap<>(), /* side-input mapping */
                  Collections.emptyList(), /* side inputs */
                  options,
                  null,
                  null,
                  DoFnSchemaInformation.create(),
                  Collections.emptyMap());
  
      DoFnOperator<String, String, String> doFnOperator = doFnOperatorSupplier.get();
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      // start a bundle
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow("regular element")));
  
      // This callback will be executed in the snapshotState function in the course of
      // finishing the currently active bundle. Everything emitted in the callback should
      // be buffered and not sent downstream.
      doFnOperator.setBundleFinishedCallback(
          () -> {
            try {
              // Clear this early for the test here because we want to finish the bundle from within
              // the callback which would otherwise cause an infinitive recursion
              doFnOperator.setBundleFinishedCallback(null);
              testHarness.processElement(
                  new StreamRecord<>(WindowedValues.valueInGlobalWindow("trigger another bundle")));
              doFnOperator.invokeFinishBundle();
              testHarness.processElement(
                  new StreamRecord<>(
                      WindowedValues.valueInGlobalWindow(
                          "check that the previous element is not flushed")));
            } catch (Exception e) {
              throw new RuntimeException(e);
            }
          });
  
      OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
  
      // Check that we have only the element which was emitted before the snapshot
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(WindowedValues.valueInGlobalWindow("regular element")));
  
      // Check that we would flush the buffered elements when continuing to run
      testHarness.processWatermark(Long.MAX_VALUE);
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("regular element"),
              WindowedValues.valueInGlobalWindow("trigger another bundle"),
              WindowedValues.valueInGlobalWindow("check that the previous element is not flushed")));
  
      testHarness.close();
  
      // Check that we would flush the buffered elements when restoring from a checkpoint
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness2 =
          new OneInputStreamOperatorTestHarness<>(doFnOperatorSupplier.get());
  
      testHarness2.initializeState(snapshot);
      testHarness2.open();
  
      testHarness2.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow("after restore")));
  
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness2.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("trigger another bundle"),
              WindowedValues.valueInGlobalWindow("check that the previous element is not flushed"),
              WindowedValues.valueInGlobalWindow("after restore")));
    }
  
    @Test
    public void testExactlyOnceBuffering() throws Exception {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(2L);
      options.setCheckpointingInterval(1L);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
      WindowedValues.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
  
      numStartBundleCalled = 0;
      DoFn<String, String> doFn =
          new DoFn<String, String>() {
            @StartBundle
            public void startBundle(StartBundleContext context) {
              numStartBundleCalled += 1;
            }
  
            @ProcessElement
            // Use RequiresStableInput to force buffering elements
            @RequiresStableInput
            public void processElement(ProcessContext context) {
              context.output(context.element());
            }
  
            @FinishBundle
            public void finishBundle(FinishBundleContext context) {
              context.output(
                  "finishBundle", BoundedWindow.TIMESTAMP_MIN_VALUE, GlobalWindow.INSTANCE);
            }
          };
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      Supplier<DoFnOperator<String, String, String>> doFnOperatorSupplier =
          () ->
              new DoFnOperator<>(
                  doFn,
                  "stepName",
                  windowedValueCoder,
                  Collections.emptyMap(),
                  outputTag,
                  Collections.emptyList(),
                  outputManagerFactory,
                  WindowingStrategy.globalDefault(),
                  new HashMap<>(), /* side-input mapping */
                  Collections.emptyList(), /* side inputs */
                  options,
                  null,
                  null,
                  DoFnSchemaInformation.create(),
                  Collections.emptyMap());
  
      DoFnOperator<String, String, String> doFnOperator = doFnOperatorSupplier.get();
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("a")));
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("b")));
  
      assertThat(Iterables.size(testHarness.getOutput()), is(0));
      assertThat(numStartBundleCalled, is(0));
  
      // create a backup and then
      OperatorSubtaskState backup = testHarness.snapshot(0, 0);
      doFnOperator.notifyCheckpointComplete(0L);
  
      assertThat(numStartBundleCalled, is(1));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("a"),
              WindowedValues.valueInGlobalWindow("b"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
  
      doFnOperator = doFnOperatorSupplier.get();
      testHarness = new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      // restore from the snapshot
      testHarness.initializeState(backup);
      testHarness.open();
  
      doFnOperator.notifyCheckpointComplete(0L);
  
      assertThat(numStartBundleCalled, is(2));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("a"),
              WindowedValues.valueInGlobalWindow("b"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
  
      // repeat to see if elements are evicted
      doFnOperator.notifyCheckpointComplete(1L);
  
      assertThat(numStartBundleCalled, is(2));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("a"),
              WindowedValues.valueInGlobalWindow("b"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
    }
  
    @Test
    public void testExactlyOnceBufferingFlushDuringDrain() throws Exception {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(2L);
      options.setCheckpointingInterval(1L);
      options.setEnableStableInputDrain(true);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
      WindowedValues.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
  
      numStartBundleCalled = 0;
      DoFn<String, String> doFn =
          new DoFn<String, String>() {
            @StartBundle
            public void startBundle(StartBundleContext context) {
              numStartBundleCalled += 1;
            }
  
            @ProcessElement
            // Use RequiresStableInput to force buffering elements
            @RequiresStableInput
            public void processElement(ProcessContext context) {
              context.output(context.element());
            }
  
            @FinishBundle
            public void finishBundle(FinishBundleContext context) {
              context.output(
                  "finishBundle", BoundedWindow.TIMESTAMP_MIN_VALUE, GlobalWindow.INSTANCE);
            }
          };
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      Supplier<DoFnOperator<String, String, String>> doFnOperatorSupplier =
          () ->
              new DoFnOperator<>(
                  doFn,
                  "stepName",
                  windowedValueCoder,
                  Collections.emptyMap(),
                  outputTag,
                  Collections.emptyList(),
                  outputManagerFactory,
                  WindowingStrategy.globalDefault(),
                  new HashMap<>(), /* side-input mapping */
                  Collections.emptyList(), /* side inputs */
                  options,
                  null,
                  null,
                  DoFnSchemaInformation.create(),
                  Collections.emptyMap());
  
      DoFnOperator<String, String, String> doFnOperator = doFnOperatorSupplier.get();
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("a")));
      testHarness.processElement(new StreamRecord<>(WindowedValues.valueInGlobalWindow("b")));
  
      assertThat(Iterables.size(testHarness.getOutput()), is(0));
      assertThat(numStartBundleCalled, is(0));
  
      // Simulate pipeline drain scenario
      OperatorSubtaskState backup = testHarness.snapshot(0, 0);
      doFnOperator.flushData();
  
      assertThat(numStartBundleCalled, is(1));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          contains(
              WindowedValues.valueInGlobalWindow("a"),
              WindowedValues.valueInGlobalWindow("b"),
              WindowedValues.valueInGlobalWindow("finishBundle")));
  
      doFnOperator = doFnOperatorSupplier.get();
      testHarness = new OneInputStreamOperatorTestHarness<>(doFnOperator);
      testHarness.open();
  
      doFnOperator.notifyCheckpointComplete(0L);
  
      assertThat(numStartBundleCalled, is(1));
      assertThat(stripStreamRecordFromWindowedValue(testHarness.getOutput()), emptyIterable());
    }
  
    @Test
    public void testExactlyOnceBufferingKeyed() throws Exception {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(2L);
      options.setCheckpointingInterval(1L);
  
      TupleTag<KV<String, String>> outputTag = new TupleTag<>("main-output");
  
      StringUtf8Coder keyCoder = StringUtf8Coder.of();
      KvToFlinkKeyKeySelector<String, String> keySelector = new KvToFlinkKeyKeySelector<>(keyCoder);
      KvCoder<String, String> kvCoder = KvCoder.of(keyCoder, StringUtf8Coder.of());
      WindowedValues.ValueOnlyWindowedValueCoder<KV<String, String>> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(kvCoder);
  
      DoFn<KV<String, String>, KV<String, String>> doFn =
          new DoFn<KV<String, String>, KV<String, String>>() {
            @StartBundle
            public void startBundle() {
              numStartBundleCalled++;
            }
  
            @ProcessElement
            // Use RequiresStableInput to force buffering elements
            @RequiresStableInput
            public void processElement(ProcessContext context) {
              context.output(context.element());
            }
  
            @FinishBundle
            public void finishBundle(FinishBundleContext context) {
              context.output(
                  KV.of("key3", "finishBundle"),
                  BoundedWindow.TIMESTAMP_MIN_VALUE,
                  GlobalWindow.INSTANCE);
            }
          };
  
      DoFnOperator.MultiOutputOutputManagerFactory<KV<String, String>> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(kvCoder, GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      Supplier<DoFnOperator<KV<String, String>, KV<String, String>, KV<String, String>>>
          doFnOperatorSupplier =
              () ->
                  new DoFnOperator<>(
                      doFn,
                      "stepName",
                      windowedValueCoder,
                      Collections.emptyMap(),
                      outputTag,
                      Collections.emptyList(),
                      outputManagerFactory,
                      WindowingStrategy.globalDefault(),
                      new HashMap<>(), /* side-input mapping */
                      Collections.emptyList(), /* side inputs */
                      options,
                      keyCoder,
                      keySelector,
                      DoFnSchemaInformation.create(),
                      Collections.emptyMap());
  
      DoFnOperator<KV<String, String>, KV<String, String>, KV<String, String>> doFnOperator =
          doFnOperatorSupplier.get();
      OneInputStreamOperatorTestHarness<
              WindowedValue<KV<String, String>>, WindowedValue<KV<String, String>>>
          testHarness =
              new KeyedOneInputStreamOperatorTestHarness<>(
                  doFnOperator, keySelector, keySelector.getProducedType());
  
      testHarness.open();
  
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key", "a"))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key", "b"))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key2", "c"))));
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow(KV.of("key2", "d"))));
  
      assertThat(
          testHarness.getOutput() + " should be empty",
          Iterables.size(testHarness.getOutput()),
          is(0));
  
      OperatorSubtaskState backup = testHarness.snapshot(0, 0);
      doFnOperator.notifyCheckpointComplete(0L);
  
      assertThat(numStartBundleCalled, is(1));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          containsInAnyOrder(
              WindowedValues.valueInGlobalWindow(KV.of("key", "a")),
              WindowedValues.valueInGlobalWindow(KV.of("key", "b")),
              WindowedValues.valueInGlobalWindow(KV.of("key2", "c")),
              WindowedValues.valueInGlobalWindow(KV.of("key2", "d")),
              WindowedValues.valueInGlobalWindow(KV.of("key3", "finishBundle"))));
  
      doFnOperator = doFnOperatorSupplier.get();
      testHarness =
          new KeyedOneInputStreamOperatorTestHarness<>(
              doFnOperator, keySelector, keySelector.getProducedType());
  
      // restore from the snapshot
      testHarness.initializeState(backup);
      testHarness.open();
  
      doFnOperator.notifyCheckpointComplete(0L);
  
      assertThat(numStartBundleCalled, is(2));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          containsInAnyOrder(
              WindowedValues.valueInGlobalWindow(KV.of("key", "a")),
              WindowedValues.valueInGlobalWindow(KV.of("key", "b")),
              WindowedValues.valueInGlobalWindow(KV.of("key2", "c")),
              WindowedValues.valueInGlobalWindow(KV.of("key2", "d")),
              WindowedValues.valueInGlobalWindow(KV.of("key3", "finishBundle"))));
  
      // repeat to see if elements are evicted
      doFnOperator.notifyCheckpointComplete(1L);
  
      assertThat(numStartBundleCalled, is(2));
      assertThat(
          stripStreamRecordFromWindowedValue(testHarness.getOutput()),
          containsInAnyOrder(
              WindowedValues.valueInGlobalWindow(KV.of("key", "a")),
              WindowedValues.valueInGlobalWindow(KV.of("key", "b")),
              WindowedValues.valueInGlobalWindow(KV.of("key2", "c")),
              WindowedValues.valueInGlobalWindow(KV.of("key2", "d")),
              WindowedValues.valueInGlobalWindow(KV.of("key3", "finishBundle"))));
    }
  
    @Test(expected = IllegalStateException.class)
    public void testFailOnRequiresStableInputAndDisabledCheckpointing() {
      TupleTag<KV<String, String>> outputTag = new TupleTag<>("main-output");
  
      StringUtf8Coder keyCoder = StringUtf8Coder.of();
      KvToFlinkKeyKeySelector<String, String> keySelector = new KvToFlinkKeyKeySelector<>(keyCoder);
      KvCoder<String, String> kvCoder = KvCoder.of(keyCoder, StringUtf8Coder.of());
      WindowedValues.ValueOnlyWindowedValueCoder<KV<String, String>> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(kvCoder);
  
      DoFn<KV<String, String>, KV<String, String>> doFn =
          new DoFn<KV<String, String>, KV<String, String>>() {
            @ProcessElement
            // Use RequiresStableInput to force buffering elements
            @RequiresStableInput
            public void processElement(ProcessContext context) {
              context.output(context.element());
            }
          };
  
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      DoFnOperator.MultiOutputOutputManagerFactory<KV<String, String>> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(kvCoder, GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      // should make the DoFnOperator creation fail
      options.setCheckpointingInterval(-1L);
      new DoFnOperator<>(
          doFn,
          "stepName",
          windowedValueCoder,
          Collections.emptyMap(),
          outputTag,
          Collections.emptyList(),
          outputManagerFactory,
          WindowingStrategy.globalDefault(),
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          options,
          keyCoder,
          keySelector,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());
    }
  
    @Test
    public void testBundleProcessingExceptionIsFatalDuringCheckpointing() throws Exception {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setMaxBundleSize(10L);
      options.setCheckpointingInterval(1L);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
  
      StringUtf8Coder coder = StringUtf8Coder.of();
      WindowedValues.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(coder);
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      DoFnOperator<String, String, String> doFnOperator =
          new DoFnOperator<>(
              new IdentityDoFn<String>() {
                @FinishBundle
                public void finishBundle() {
                  throw new RuntimeException("something went wrong here");
                }
              },
              "stepName",
              windowedValueCoder,
              Collections.emptyMap(),
              outputTag,
              Collections.emptyList(),
              outputManagerFactory,
              WindowingStrategy.globalDefault(),
              new HashMap<>(), /* side-input mapping */
              Collections.emptyList(), /* side inputs */
              options,
              null,
              null,
              DoFnSchemaInformation.create(),
              Collections.emptyMap());
  
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      // start a bundle
      testHarness.processElement(
          new StreamRecord<>(WindowedValues.valueInGlobalWindow("regular element")));
  
      // Make sure we throw Error, not a regular Exception.
      // A regular exception would just cause the checkpoint to fail.
      assertThrows(Error.class, () -> testHarness.snapshot(0, 0));
    }
  
    @Test
    public void testAccumulatorRegistrationOnOperatorClose() throws Exception {
      DoFnOperator<String, String, String> doFnOperator = getOperatorForCleanupInspection();
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(doFnOperator);
  
      testHarness.open();
  
      String metricContainerFieldName = "flinkMetricContainer";
      FlinkMetricContainer monitoredContainer =
          Mockito.spy(
              (FlinkMetricContainer)
                  Whitebox.getInternalState(doFnOperator, metricContainerFieldName));
      Whitebox.setInternalState(doFnOperator, metricContainerFieldName, monitoredContainer);
  
      // Closes and disposes the operator
      testHarness.close();
      // Ensure that dispose has the metrics code
      doFnOperator.cleanUp();
      Mockito.verify(monitoredContainer, Mockito.times(2)).registerMetricsForPipelineResult();
    }
  
    /**
     * Ensures Jackson cache is cleaned to get rid of any references to the Flink Classloader. See
     * https://jira.apache.org/jira/browse/BEAM-6460
     */
    @Test
    public void testRemoveCachedClassReferences() throws Exception {
      OneInputStreamOperatorTestHarness<WindowedValue<String>, WindowedValue<String>> testHarness =
          new OneInputStreamOperatorTestHarness<>(getOperatorForCleanupInspection());
  
      LRUMap<?, ?> typeCache = Whitebox.getInternalState(TypeFactory.defaultInstance(), "_typeCache");
      assertThat(typeCache.size(), greaterThan(0));
      testHarness.open();
      testHarness.close();
      assertThat(typeCache.size(), is(0));
    }
  
    private static DoFnOperator<String, String, String> getOperatorForCleanupInspection() {
      FlinkPipelineOptions options = FlinkPipelineOptions.defaults();
      options.setParallelism(4);
  
      TupleTag<String> outputTag = new TupleTag<>("main-output");
      WindowedValues.ValueOnlyWindowedValueCoder<String> windowedValueCoder =
          WindowedValues.getValueOnlyCoder(StringUtf8Coder.of());
      IdentityDoFn<String> doFn =
          new IdentityDoFn<String>() {
            @FinishBundle
            public void finishBundle(FinishBundleContext context) {
              context.output(
                  "finishBundle", BoundedWindow.TIMESTAMP_MIN_VALUE, GlobalWindow.INSTANCE);
            }
          };
  
      DoFnOperator.MultiOutputOutputManagerFactory<String> outputManagerFactory =
          new DoFnOperator.MultiOutputOutputManagerFactory<>(
              outputTag,
              WindowedValues.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE),
              new SerializablePipelineOptions(options));
  
      return new DoFnOperator<>(
          doFn,
          "stepName",
          windowedValueCoder,
          Collections.emptyMap(),
          outputTag,
          Collections.emptyList(),
          outputManagerFactory,
          WindowingStrategy.globalDefault(),
          new HashMap<>(), /* side-input mapping */
          Collections.emptyList(), /* side inputs */
          options,
          null,
          null,
          DoFnSchemaInformation.create(),
          Collections.emptyMap());
    }
  
    private Iterable<WindowedValue<String>> stripStreamRecord(Iterable<?> input) {
      return FluentIterable.from(input)
          .filter(o -> o instanceof StreamRecord)
          .transform(
              new Function<Object, WindowedValue<String>>() {
                @Override
                @SuppressWarnings({"unchecked", "rawtypes"})
                public @Nullable WindowedValue<String> apply(@Nullable Object o) {
                  if (o instanceof StreamRecord) {
                    return (WindowedValue<String>) ((StreamRecord) o).getValue();
                  }
                  throw new RuntimeException("unreachable");
                }
              });
    }
  
    private static class MultiOutputDoFn extends DoFn<String, String> {
      private final TupleTag<String> additionalOutput1;
      private final TupleTag<String> additionalOutput2;
  
      public MultiOutputDoFn(TupleTag<String> additionalOutput1, TupleTag<String> additionalOutput2) {
        this.additionalOutput1 = additionalOutput1;
        this.additionalOutput2 = additionalOutput2;
      }
  
      @ProcessElement
      public void processElement(ProcessContext c) {
        if ("one".equals(c.element())) {
          c.output(additionalOutput1, "extra: one");
        } else if ("two".equals(c.element())) {
          c.output(additionalOutput2, "extra: two");
        } else {
          c.output("got: " + c.element());
          c.output(additionalOutput1, "got: " + c.element());
          c.output(additionalOutput2, "got: " + c.element());
        }
      }
    }
  
    private static class IdentityDoFn<T> extends DoFn<T, T> {
  
      @ProcessElement
      public void processElement(ProcessContext c) {
        c.output(c.element());
      }
    }
  
    private WindowedValue<Iterable<?>> valuesInWindow(
        Iterable<?> values, Instant timestamp, BoundedWindow window) {
      return WindowedValues.of(values, timestamp, window, PaneInfo.NO_FIRING);
    }
  
    private <T> WindowedValue<T> valueInWindow(T value, Instant timestamp, BoundedWindow window) {
      return WindowedValues.of(value, timestamp, window, PaneInfo.NO_FIRING);
    }
  
    private interface TestHarnessFactory<T> {
      T create() throws Exception;
    }
  }
  