/*
 * 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.core.construction.PTransformTranslation.PAR_DO_TRANSFORM_URN;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.hasSize;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.iterableWithSize;
import static org.hamcrest.collection.IsIterableContainingInOrder.contains;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyString;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.Lock;
import javax.annotation.Nullable;
import org.apache.beam.model.fnexecution.v1.BeamFnApi;
import org.apache.beam.model.pipeline.v1.RunnerApi;
import org.apache.beam.model.pipeline.v1.RunnerApi.Components;
import org.apache.beam.model.pipeline.v1.RunnerApi.ExecutableStagePayload;
import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection;
import org.apache.beam.runners.core.InMemoryStateInternals;
import org.apache.beam.runners.core.InMemoryTimerInternals;
import org.apache.beam.runners.core.StateNamespace;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.core.StateTags;
import org.apache.beam.runners.core.StatefulDoFnRunner;
import org.apache.beam.runners.core.TimerInternals;
import org.apache.beam.runners.flink.FlinkPipelineOptions;
import org.apache.beam.runners.flink.metrics.DoFnRunnerWithMetricsUpdate;
import org.apache.beam.runners.flink.streaming.FlinkStateInternalsTest;
import org.apache.beam.runners.flink.translation.functions.FlinkExecutableStageContextFactory;
import org.apache.beam.runners.flink.translation.types.CoderTypeInformation;
import org.apache.beam.runners.flink.translation.utils.NoopLock;
import org.apache.beam.runners.fnexecution.control.BundleProgressHandler;
import org.apache.beam.runners.fnexecution.control.ExecutableStageContext;
import org.apache.beam.runners.fnexecution.control.OutputReceiverFactory;
import org.apache.beam.runners.fnexecution.control.ProcessBundleDescriptors;
import org.apache.beam.runners.fnexecution.control.RemoteBundle;
import org.apache.beam.runners.fnexecution.control.StageBundleFactory;
import org.apache.beam.runners.fnexecution.provisioning.JobInfo;
import org.apache.beam.runners.fnexecution.state.StateRequestHandler;
import org.apache.beam.runners.fnexecution.state.StateRequestHandlers;
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.VoidCoder;
import org.apache.beam.sdk.fn.IdGenerator;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.TimeDomain;
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.util.CoderUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.WindowingStrategy;
import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.ByteString;
import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.Struct;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Charsets;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap;
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.apache.beam.vendor.sdk.v2.sdk.extensions.protobuf.ByteStringCoder;
import org.apache.commons.lang3.SerializationUtils;
import org.apache.commons.lang3.mutable.MutableObject;
import org.apache.flink.api.common.cache.DistributedCache;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.runtime.state.KeyedStateBackend;
import org.apache.flink.streaming.api.watermark.Watermark;
import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
import org.apache.flink.streaming.util.KeyedOneInputStreamOperatorTestHarness;
import org.apache.flink.streaming.util.OneInputStreamOperatorTestHarness;
import org.apache.flink.util.OutputTag;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;
import org.powermock.reflect.Whitebox;

/** Tests for {@link ExecutableStageDoFnOperator}. */
@RunWith(JUnit4.class)
public class ExecutableStageDoFnOperatorTest {
  @Rule public ExpectedException thrown = ExpectedException.none();

  @Mock private RuntimeContext runtimeContext;
  @Mock private DistributedCache distributedCache;
  @Mock private ExecutableStageContext stageContext;
  @Mock private StageBundleFactory stageBundleFactory;
  @Mock private StateRequestHandler stateRequestHandler;
  @Mock private ProcessBundleDescriptors.ExecutableProcessBundleDescriptor processBundleDescriptor;

  // NOTE: ExecutableStage.fromPayload expects exactly one input, so we provide one here. These unit
  // tests in general ignore the executable stage itself and mock around it.
  private final ExecutableStagePayload stagePayload =
      ExecutableStagePayload.newBuilder()
          .setInput("input")
          .setComponents(
              Components.newBuilder()
                  .putPcollections("input", PCollection.getDefaultInstance())
                  .build())
          .build();

  private final String stateId = "userState";
  private final ExecutableStagePayload stagePayloadWithUserState =
      stagePayload
          .toBuilder()
          .setComponents(
              stagePayload
                  .getComponents()
                  .toBuilder()
                  .putTransforms(
                      "transform",
                      RunnerApi.PTransform.newBuilder()
                          .setSpec(
                              RunnerApi.FunctionSpec.newBuilder()
                                  .setUrn(PAR_DO_TRANSFORM_URN)
                                  .build())
                          .putInputs("input", "input")
                          .build())
                  .build())
          .addUserStates(
              ExecutableStagePayload.UserStateId.newBuilder()
                  .setLocalName(stateId)
                  .setTransformId("transform")
                  .build())
          .build();

  private final JobInfo jobInfo =
      JobInfo.create("job-id", "job-name", "retrieval-token", Struct.getDefaultInstance());

  @Before
  public void setUpMocks() {
    MockitoAnnotations.initMocks(this);
    when(runtimeContext.getDistributedCache()).thenReturn(distributedCache);
    when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);
    when(processBundleDescriptor.getTimerSpecs()).thenReturn(Collections.emptyMap());
    when(processBundleDescriptor.getBagUserStateSpecs()).thenReturn(Collections.emptyMap());
    when(stageBundleFactory.getProcessBundleDescriptor()).thenReturn(processBundleDescriptor);
  }

  @Test
  public void sdkErrorsSurfaceOnClose() throws Exception {
    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
    ExecutableStageDoFnOperator<Integer, Integer> operator =
        getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

    OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
        new OneInputStreamOperatorTestHarness<>(operator);

    testHarness.open();

    @SuppressWarnings("unchecked")
    RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
    when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);

    @SuppressWarnings("unchecked")
    FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
    when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

    Exception expected = new RuntimeException(new Exception());
    doThrow(expected).when(bundle).close();
    thrown.expectCause(is(expected));

    operator.processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(0)));
    testHarness.close();
  }

  @Test
  public void expectedInputsAreSent() throws Exception {
    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
    ExecutableStageDoFnOperator<Integer, Integer> operator =
        getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

    @SuppressWarnings("unchecked")
    RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
    when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);

    @SuppressWarnings("unchecked")
    FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
    when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

    WindowedValue<Integer> one = WindowedValue.valueInGlobalWindow(1);
    WindowedValue<Integer> two = WindowedValue.valueInGlobalWindow(2);
    WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);

    OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
        new OneInputStreamOperatorTestHarness<>(operator);

    testHarness.open();

    testHarness.processElement(new StreamRecord<>(one));
    testHarness.processElement(new StreamRecord<>(two));
    testHarness.processElement(new StreamRecord<>(three));

    verify(receiver).accept(one);
    verify(receiver).accept(two);
    verify(receiver).accept(three);
    verifyNoMoreInteractions(receiver);

    testHarness.close();
  }

  @Test
  public void outputsAreTaggedCorrectly() throws Exception {

    WindowedValue.ValueOnlyWindowedValueCoder<Integer> coder =
        WindowedValue.getValueOnlyCoder(VarIntCoder.of());

    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    TupleTag<Integer> additionalOutput1 = new TupleTag<>("output-1");
    TupleTag<Integer> additionalOutput2 = new TupleTag<>("output-2");
    ImmutableMap<TupleTag<?>, OutputTag<?>> tagsToOutputTags =
        ImmutableMap.<TupleTag<?>, OutputTag<?>>builder()
            .put(additionalOutput1, new OutputTag<String>(additionalOutput1.getId()) {})
            .put(additionalOutput2, new OutputTag<String>(additionalOutput2.getId()) {})
            .build();
    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.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(
            mainOutput, tagsToOutputTags, tagsToCoders, tagsToIds);

    WindowedValue<Integer> zero = WindowedValue.valueInGlobalWindow(0);
    WindowedValue<Integer> three = WindowedValue.valueInGlobalWindow(3);
    WindowedValue<Integer> four = WindowedValue.valueInGlobalWindow(4);
    WindowedValue<Integer> five = WindowedValue.valueInGlobalWindow(5);

    // We use a real StageBundleFactory here in order to exercise the output receiver factory.
    StageBundleFactory stageBundleFactory =
        new StageBundleFactory() {

          private boolean onceEmitted;

          @Override
          public RemoteBundle getBundle(
              OutputReceiverFactory receiverFactory,
              StateRequestHandler stateRequestHandler,
              BundleProgressHandler progressHandler) {
            return new RemoteBundle() {
              @Override
              public String getId() {
                return "bundle-id";
              }

              @Override
              public Map<String, FnDataReceiver> getInputReceivers() {
                return ImmutableMap.of(
                    "input",
                    input -> {
                      /* Ignore input*/
                    });
              }

              @Override
              public void close() throws Exception {
                if (onceEmitted) {
                  return;
                }
                // Emit all values to the runner when the bundle is closed.
                receiverFactory.create(mainOutput.getId()).accept(three);
                receiverFactory.create(additionalOutput1.getId()).accept(four);
                receiverFactory.create(additionalOutput2.getId()).accept(five);
                onceEmitted = true;
              }
            };
          }

          @Override
          public ProcessBundleDescriptors.ExecutableProcessBundleDescriptor
              getProcessBundleDescriptor() {
            return processBundleDescriptor;
          }

          @Override
          public void close() {}
        };
    // Wire the stage bundle factory into our context.
    when(stageContext.getStageBundleFactory(any())).thenReturn(stageBundleFactory);

    ExecutableStageDoFnOperator<Integer, Integer> operator =
        getOperator(
            mainOutput,
            ImmutableList.of(additionalOutput1, additionalOutput2),
            outputManagerFactory);

    OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
        new OneInputStreamOperatorTestHarness<>(operator);

    long watermark = testHarness.getCurrentWatermark() + 1;
    testHarness.open();

    testHarness.processElement(new StreamRecord<>(zero));

    testHarness.processWatermark(watermark);
    watermark++;
    testHarness.processWatermark(watermark);

    assertEquals(watermark, testHarness.getCurrentWatermark());
    // watermark hold until bundle complete
    assertEquals(0, testHarness.getOutput().size());

    testHarness.close(); // triggers finish bundle

    assertThat(
        testHarness.getOutput(),
        contains(
            new StreamRecord<>(three), new Watermark(watermark), new Watermark(Long.MAX_VALUE)));

    assertThat(
        testHarness.getSideOutput(tagsToOutputTags.get(additionalOutput1)),
        contains(new StreamRecord<>(four)));

    assertThat(
        testHarness.getSideOutput(tagsToOutputTags.get(additionalOutput2)),
        contains(new StreamRecord<>(five)));
  }

  @Test
  public void testStageBundleClosed() throws Exception {
    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
    ExecutableStageDoFnOperator<Integer, Integer> operator =
        getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

    OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
        new OneInputStreamOperatorTestHarness<>(operator);

    RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
    when(bundle.getInputReceivers())
        .thenReturn(
            ImmutableMap.<String, FnDataReceiver<WindowedValue>>builder()
                .put("input", Mockito.mock(FnDataReceiver.class))
                .build());
    when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);

    testHarness.open();
    testHarness.close();

    verify(stageBundleFactory).getProcessBundleDescriptor();
    verify(stageBundleFactory).close();
    verify(stageContext).close();
    verifyNoMoreInteractions(stageBundleFactory);

    // close() will also call dispose(), but call again to verify no new bundle
    // is created afterwards
    operator.dispose();
    verifyNoMoreInteractions(bundle);
  }

  @Test
  @SuppressWarnings("unchecked")
  public void testEnsureStateCleanupWithKeyedInput() throws Exception {
    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VarIntCoder.of());
    VarIntCoder keyCoder = VarIntCoder.of();
    ExecutableStageDoFnOperator<Integer, Integer> operator =
        getOperator(
            mainOutput,
            Collections.emptyList(),
            outputManagerFactory,
            WindowingStrategy.globalDefault(),
            keyCoder,
            WindowedValue.getFullCoder(keyCoder, GlobalWindow.Coder.INSTANCE));

    KeyedOneInputStreamOperatorTestHarness<Integer, WindowedValue<Integer>, WindowedValue<Integer>>
        testHarness =
            new KeyedOneInputStreamOperatorTestHarness(
                operator, val -> val, new CoderTypeInformation<>(keyCoder));

    RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
    when(bundle.getInputReceivers())
        .thenReturn(
            ImmutableMap.<String, FnDataReceiver<WindowedValue>>builder()
                .put("input", Mockito.mock(FnDataReceiver.class))
                .build());
    when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);

    testHarness.open();

    Object doFnRunner = Whitebox.getInternalState(operator, "doFnRunner");
    assertThat(doFnRunner, instanceOf(DoFnRunnerWithMetricsUpdate.class));

    // There should be a StatefulDoFnRunner installed which takes care of clearing state
    Object statefulDoFnRunner = Whitebox.getInternalState(doFnRunner, "delegate");
    assertThat(statefulDoFnRunner, instanceOf(StatefulDoFnRunner.class));
  }

  @Test
  public void testEnsureStateCleanupWithKeyedInputCleanupTimer() {
    InMemoryTimerInternals inMemoryTimerInternals = new InMemoryTimerInternals();
    KeyedStateBackend keyedStateBackend = Mockito.mock(KeyedStateBackend.class);
    Lock stateBackendLock = Mockito.mock(Lock.class);
    StringUtf8Coder keyCoder = StringUtf8Coder.of();
    GlobalWindow window = GlobalWindow.INSTANCE;
    GlobalWindow.Coder windowCoder = GlobalWindow.Coder.INSTANCE;

    // Test that cleanup timer is set correctly
    ExecutableStageDoFnOperator.CleanupTimer cleanupTimer =
        new ExecutableStageDoFnOperator.CleanupTimer<>(
            inMemoryTimerInternals,
            stateBackendLock,
            WindowingStrategy.globalDefault(),
            keyCoder,
            windowCoder,
            keyedStateBackend);
    cleanupTimer.setForWindow(KV.of("key", "string"), window);

    Mockito.verify(stateBackendLock).lock();
    ByteBuffer key = FlinkKeyUtils.encodeKey("key", keyCoder);
    Mockito.verify(keyedStateBackend).setCurrentKey(key);
    assertThat(
        inMemoryTimerInternals.getNextTimer(TimeDomain.EVENT_TIME),
        is(window.maxTimestamp().plus(1)));
    Mockito.verify(stateBackendLock).unlock();
  }

  @Test
  public void testEnsureStateCleanupWithKeyedInputStateCleaner() {
    GlobalWindow.Coder windowCoder = GlobalWindow.Coder.INSTANCE;
    InMemoryStateInternals<String> stateInternals = InMemoryStateInternals.forKey("key");
    List<String> userStateNames = ImmutableList.of("state1", "state2");
    ImmutableList.Builder<BagState<String>> bagStateBuilder = ImmutableList.builder();
    for (String userStateName : userStateNames) {
      BagState<String> state =
          stateInternals.state(
              StateNamespaces.window(windowCoder, GlobalWindow.INSTANCE),
              StateTags.bag(userStateName, StringUtf8Coder.of()));
      bagStateBuilder.add(state);
      state.add("this should be cleaned");
    }
    ImmutableList<BagState<String>> bagStates = bagStateBuilder.build();

    MutableObject<ByteBuffer> key =
        new MutableObject<>(
            ByteBuffer.wrap(stateInternals.getKey().getBytes(StandardCharsets.UTF_8)));

    // Test that state is cleaned up correctly
    ExecutableStageDoFnOperator.StateCleaner stateCleaner =
        new ExecutableStageDoFnOperator.StateCleaner(
            userStateNames, windowCoder, () -> key.getValue());
    for (BagState<String> bagState : bagStates) {
      assertThat(Iterables.size(bagState.read()), is(1));
    }

    stateCleaner.clearForWindow(GlobalWindow.INSTANCE);
    stateCleaner.cleanupState(stateInternals, (k) -> key.setValue(k));

    for (BagState<String> bagState : bagStates) {
      assertThat(Iterables.size(bagState.read()), is(0));
    }
  }

  @Test
  public void testEnsureDeferredStateCleanupTimerFiring() throws Exception {
    testEnsureDeferredStateCleanupTimerFiring(false);
    testEnsureDeferredStateCleanupTimerFiring(true);
  }

  private void testEnsureDeferredStateCleanupTimerFiring(boolean withCheckpointing)
      throws Exception {
    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
    StringUtf8Coder keyCoder = StringUtf8Coder.of();

    WindowingStrategy windowingStrategy =
        WindowingStrategy.of(FixedWindows.of(Duration.millis(1000)));

    KvCoder<String, Integer> kvCoder = KvCoder.of(keyCoder, VarIntCoder.of());
    ExecutableStageDoFnOperator<Integer, Integer> operator =
        getOperator(
            mainOutput,
            Collections.emptyList(),
            outputManagerFactory,
            windowingStrategy,
            keyCoder,
            WindowedValue.getFullCoder(kvCoder, windowingStrategy.getWindowFn().windowCoder()));

    @SuppressWarnings("unchecked")
    RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
    when(stageBundleFactory.getBundle(any(), any(), any())).thenReturn(bundle);

    String timerInputId = "timerInput";
    AtomicBoolean timerInputReceived = new AtomicBoolean();
    IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(1000));
    IntervalWindow.IntervalWindowCoder windowCoder = IntervalWindow.IntervalWindowCoder.of();
    WindowedValue<KV<String, Integer>> one =
        WindowedValue.of(
            KV.of("one", 1), window.maxTimestamp(), ImmutableList.of(window), PaneInfo.NO_FIRING);

    FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
    FnDataReceiver<WindowedValue<?>> timerReceiver = Mockito.mock(FnDataReceiver.class);
    doAnswer(
            (invocation) -> {
              timerInputReceived.set(true);
              return null;
            })
        .when(timerReceiver)
        .accept(any());

    when(bundle.getInputReceivers())
        .thenReturn(ImmutableMap.of("input", receiver, timerInputId, timerReceiver));

    KeyedOneInputStreamOperatorTestHarness<
            String, WindowedValue<KV<String, Integer>>, WindowedValue<KV<String, Integer>>>
        testHarness =
            new KeyedOneInputStreamOperatorTestHarness(
                operator, val -> val, new CoderTypeInformation<>(keyCoder));

    testHarness.open();

    Lock stateBackendLock = (Lock) Whitebox.getInternalState(operator, "stateBackendLock");
    stateBackendLock.lock();

    KeyedStateBackend<ByteBuffer> keyedStateBackend = operator.getKeyedStateBackend();
    ByteBuffer key = FlinkKeyUtils.encodeKey(one.getValue().getKey(), keyCoder);
    keyedStateBackend.setCurrentKey(key);

    DoFnOperator.FlinkTimerInternals timerInternals =
        (DoFnOperator.FlinkTimerInternals) Whitebox.getInternalState(operator, "timerInternals");

    Object doFnRunner = Whitebox.getInternalState(operator, "doFnRunner");
    Object delegate = Whitebox.getInternalState(doFnRunner, "delegate");
    Object stateCleaner = Whitebox.getInternalState(delegate, "stateCleaner");
    Collection<?> cleanupTimers =
        (Collection) Whitebox.getInternalState(stateCleaner, "cleanupQueue");

    // create some state which can be cleaned up
    assertThat(testHarness.numKeyedStateEntries(), is(0));
    StateNamespace stateNamespace = StateNamespaces.window(windowCoder, window);
    BagState<String> state =
        operator.keyedStateInternals.state(
            stateNamespace, StateTags.bag(stateId, StringUtf8Coder.of()));
    state.add("testUserState");
    assertThat(testHarness.numKeyedStateEntries(), is(1));

    // user timer that fires after the end of the window and after state cleanup
    TimerInternals.TimerData userTimer =
        TimerInternals.TimerData.of(
            timerInputId, stateNamespace, window.maxTimestamp().plus(1), TimeDomain.EVENT_TIME);
    timerInternals.setTimer(userTimer);

    // start of bundle
    testHarness.processElement(new StreamRecord<>(one));
    verify(receiver).accept(one);

    // move watermark past cleanup and user timer while bundle in progress
    operator.processWatermark(new Watermark(window.maxTimestamp().plus(2).getMillis()));

    // due to watermark hold the timers won't fire at this point
    assertFalse("Watermark must be held back until bundle is complete.", timerInputReceived.get());
    assertThat(cleanupTimers, hasSize(0));

    if (withCheckpointing) {
      // Upon checkpointing, the bundle is finished and the watermark advances;
      // timers can fire. Note: The bundle is ensured to be finished.
      testHarness.snapshot(0, 0);

      // The user timer was scheduled to fire after cleanup, but executes first
      assertTrue("Timer should have been triggered.", timerInputReceived.get());
      // Cleanup will be executed after the bundle is complete
      assertThat(cleanupTimers, hasSize(0));
      verifyNoMoreInteractions(receiver);
    } else {
      // Upon finishing a bundle, the watermark advances; timers can fire.
      // Note that this will finish the current bundle, but will also start a new one
      // when timers fire as part of advancing the watermark
      operator.invokeFinishBundle();

      // The user timer was scheduled to fire after cleanup, but executes first
      assertTrue("Timer should have been triggered.", timerInputReceived.get());
      // Cleanup will be executed after the bundle is complete
      assertThat(cleanupTimers, hasSize(1));
      verifyNoMoreInteractions(receiver);

      // Finish bundle which has been started by finishing the bundle
      operator.invokeFinishBundle();
      assertThat(cleanupTimers, hasSize(0));
    }

    assertThat(testHarness.numKeyedStateEntries(), is(0));

    testHarness.close();
  }

  @Test
  public void testCacheTokenHandling() throws Exception {
    InMemoryStateInternals test = InMemoryStateInternals.forKey("test");
    KeyedStateBackend<ByteBuffer> stateBackend = FlinkStateInternalsTest.createStateBackend();

    // User state the cache token is valid for the lifetime of the operator
    for (String expectedToken : new String[] {"first token", "second token"}) {
      final IdGenerator cacheTokenGenerator = () -> expectedToken;
      ExecutableStageDoFnOperator.BagUserStateFactory<ByteString, Integer, GlobalWindow>
          bagUserStateFactory =
              new ExecutableStageDoFnOperator.BagUserStateFactory<>(
                  cacheTokenGenerator, test, stateBackend, NoopLock.get());

      ByteString key1 = ByteString.copyFrom("key1", Charsets.UTF_8);
      ByteString key2 = ByteString.copyFrom("key2", Charsets.UTF_8);

      Map<String, Map<String, ProcessBundleDescriptors.BagUserStateSpec>> userStateMapMock =
          Mockito.mock(Map.class);
      Map<String, ProcessBundleDescriptors.BagUserStateSpec> transformMap = Mockito.mock(Map.class);

      final String userState1 = "userstate1";
      ProcessBundleDescriptors.BagUserStateSpec bagUserStateSpec1 = mockBagUserState(userState1);
      when(transformMap.get(userState1)).thenReturn(bagUserStateSpec1);

      final String userState2 = "userstate2";
      ProcessBundleDescriptors.BagUserStateSpec bagUserStateSpec2 = mockBagUserState(userState2);
      when(transformMap.get(userState2)).thenReturn(bagUserStateSpec2);

      when(userStateMapMock.get(anyString())).thenReturn(transformMap);
      when(processBundleDescriptor.getBagUserStateSpecs()).thenReturn(userStateMapMock);
      StateRequestHandler stateRequestHandler =
          StateRequestHandlers.forBagUserStateHandlerFactory(
              processBundleDescriptor, bagUserStateFactory);

      // There should be no cache token available before any requests have been made
      assertThat(stateRequestHandler.getCacheTokens(), iterableWithSize(0));

      // Make a request to generate initial cache token
      stateRequestHandler.handle(getRequest(key1, userState1));
      BeamFnApi.ProcessBundleRequest.CacheToken cacheTokenStruct =
          Iterables.getOnlyElement(stateRequestHandler.getCacheTokens());
      assertThat(cacheTokenStruct.hasUserState(), is(true));
      ByteString cacheToken = cacheTokenStruct.getToken();
      final ByteString expectedCacheToken =
          ByteString.copyFrom(expectedToken.getBytes(Charsets.UTF_8));
      assertThat(cacheToken, is(expectedCacheToken));

      List<RequestGenerator> generators =
          Arrays.asList(
              ExecutableStageDoFnOperatorTest::getRequest,
              ExecutableStageDoFnOperatorTest::getAppend,
              ExecutableStageDoFnOperatorTest::getClear);

      for (RequestGenerator req : generators) {
        // For every state read the tokens remains unchanged
        stateRequestHandler.handle(req.makeRequest(key1, userState1));
        assertThat(
            Iterables.getOnlyElement(stateRequestHandler.getCacheTokens()).getToken(),
            is(expectedCacheToken));

        // The token is still valid for another key in the same key range
        stateRequestHandler.handle(req.makeRequest(key2, userState1));
        assertThat(
            Iterables.getOnlyElement(stateRequestHandler.getCacheTokens()).getToken(),
            is(expectedCacheToken));

        // The token is still valid for another state cell in the same key range
        stateRequestHandler.handle(req.makeRequest(key2, userState2));
        assertThat(
            Iterables.getOnlyElement(stateRequestHandler.getCacheTokens()).getToken(),
            is(expectedCacheToken));
      }
    }
  }

  private interface RequestGenerator {
    BeamFnApi.StateRequest makeRequest(ByteString key, String userStateId) throws Exception;
  }

  private static BeamFnApi.StateRequest getRequest(ByteString key, String userStateId)
      throws Exception {
    BeamFnApi.StateRequest.Builder builder = stateRequest(key, userStateId);
    builder.setGet(BeamFnApi.StateGetRequest.newBuilder().build());
    return builder.build();
  }

  private static BeamFnApi.StateRequest getAppend(ByteString key, String userStateId)
      throws Exception {
    BeamFnApi.StateRequest.Builder builder = stateRequest(key, userStateId);
    builder.setAppend(BeamFnApi.StateAppendRequest.newBuilder().build());
    return builder.build();
  }

  private static BeamFnApi.StateRequest getClear(ByteString key, String userStateId)
      throws Exception {
    BeamFnApi.StateRequest.Builder builder = stateRequest(key, userStateId);
    builder.setClear(BeamFnApi.StateClearRequest.newBuilder().build());
    return builder.build();
  }

  private static BeamFnApi.StateRequest.Builder stateRequest(ByteString key, String userStateId)
      throws Exception {
    return BeamFnApi.StateRequest.newBuilder()
        .setStateKey(
            BeamFnApi.StateKey.newBuilder()
                .setBagUserState(
                    BeamFnApi.StateKey.BagUserState.newBuilder()
                        .setTransformId("transform")
                        .setKey(key)
                        .setUserStateId(userStateId)
                        .setWindow(
                            ByteString.copyFrom(
                                CoderUtils.encodeToByteArray(
                                    GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE)))
                        .build()));
  }

  private static ProcessBundleDescriptors.BagUserStateSpec mockBagUserState(String userStateId) {
    ProcessBundleDescriptors.BagUserStateSpec bagUserStateMock =
        Mockito.mock(ProcessBundleDescriptors.BagUserStateSpec.class);
    when(bagUserStateMock.keyCoder()).thenReturn(ByteStringCoder.of());
    when(bagUserStateMock.valueCoder()).thenReturn(ByteStringCoder.of());
    when(bagUserStateMock.transformId()).thenReturn("transformId");
    when(bagUserStateMock.userStateId()).thenReturn(userStateId);
    when(bagUserStateMock.windowCoder()).thenReturn(GlobalWindow.Coder.INSTANCE);
    return bagUserStateMock;
  }

  @Test
  public void testSerialization() {
    WindowedValue.ValueOnlyWindowedValueCoder<Integer> coder =
        WindowedValue.getValueOnlyCoder(VarIntCoder.of());

    TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
    TupleTag<Integer> additionalOutput = new TupleTag<>("additional-output");
    ImmutableMap<TupleTag<?>, OutputTag<?>> tagsToOutputTags =
        ImmutableMap.<TupleTag<?>, OutputTag<?>>builder()
            .put(
                additionalOutput,
                new OutputTag<>(additionalOutput.getId(), TypeInformation.of(Integer.class)))
            .build();
    ImmutableMap<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders =
        ImmutableMap.<TupleTag<?>, Coder<WindowedValue<?>>>builder()
            .put(mainOutput, (Coder) coder)
            .put(additionalOutput, coder)
            .build();
    ImmutableMap<TupleTag<?>, Integer> tagsToIds =
        ImmutableMap.<TupleTag<?>, Integer>builder()
            .put(mainOutput, 0)
            .put(additionalOutput, 1)
            .build();

    DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
        new DoFnOperator.MultiOutputOutputManagerFactory(
            mainOutput, tagsToOutputTags, tagsToCoders, tagsToIds);

    FlinkPipelineOptions options = PipelineOptionsFactory.as(FlinkPipelineOptions.class);

    ExecutableStageDoFnOperator<Integer, Integer> operator =
        new ExecutableStageDoFnOperator<>(
            "transform",
            null,
            null,
            Collections.emptyMap(),
            mainOutput,
            ImmutableList.of(additionalOutput),
            outputManagerFactory,
            Collections.emptyMap() /* sideInputTagMapping */,
            Collections.emptyList() /* sideInputs */,
            Collections.emptyMap() /* sideInputId mapping */,
            options,
            stagePayload,
            jobInfo,
            FlinkExecutableStageContextFactory.getInstance(),
            createOutputMap(mainOutput, ImmutableList.of(additionalOutput)),
            WindowingStrategy.globalDefault(),
            null,
            null);

    ExecutableStageDoFnOperator<Integer, Integer> clone = SerializationUtils.clone(operator);
    assertNotNull(clone);
    assertNotEquals(operator, clone);
  }

  /**
   * Creates a {@link ExecutableStageDoFnOperator}. Sets the runtime context to {@link
   * #runtimeContext}. The context factory is mocked to return {@link #stageContext} every time. The
   * behavior of the stage context itself is unchanged.
   */
  private ExecutableStageDoFnOperator<Integer, Integer> getOperator(
      TupleTag<Integer> mainOutput,
      List<TupleTag<?>> additionalOutputs,
      DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory) {
    return getOperator(
        mainOutput,
        additionalOutputs,
        outputManagerFactory,
        WindowingStrategy.globalDefault(),
        null,
        WindowedValue.getFullCoder(StringUtf8Coder.of(), GlobalWindow.Coder.INSTANCE));
  }

  private ExecutableStageDoFnOperator<Integer, Integer> getOperator(
      TupleTag<Integer> mainOutput,
      List<TupleTag<?>> additionalOutputs,
      DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory,
      WindowingStrategy windowingStrategy,
      @Nullable Coder keyCoder,
      @Nullable Coder windowedInputCoder) {

    FlinkExecutableStageContextFactory contextFactory =
        Mockito.mock(FlinkExecutableStageContextFactory.class);
    when(contextFactory.get(any())).thenReturn(stageContext);

    final ExecutableStagePayload stagePayload;
    if (keyCoder != null) {
      stagePayload = this.stagePayloadWithUserState;
    } else {
      stagePayload = this.stagePayload;
    }

    ExecutableStageDoFnOperator<Integer, Integer> operator =
        new ExecutableStageDoFnOperator<>(
            "transform",
            windowedInputCoder,
            null,
            Collections.emptyMap(),
            mainOutput,
            additionalOutputs,
            outputManagerFactory,
            Collections.emptyMap() /* sideInputTagMapping */,
            Collections.emptyList() /* sideInputs */,
            Collections.emptyMap() /* sideInputId mapping */,
            PipelineOptionsFactory.as(FlinkPipelineOptions.class),
            stagePayload,
            jobInfo,
            contextFactory,
            createOutputMap(mainOutput, additionalOutputs),
            windowingStrategy,
            keyCoder,
            keyCoder != null ? new KvToByteBufferKeySelector<>(keyCoder) : null);

    Whitebox.setInternalState(operator, "stateRequestHandler", stateRequestHandler);
    return operator;
  }

  private static Map<String, TupleTag<?>> createOutputMap(
      TupleTag mainOutput, List<TupleTag<?>> additionalOutputs) {
    Map<String, TupleTag<?>> outputMap = new HashMap<>(additionalOutputs.size() + 1);
    if (mainOutput != null) {
      outputMap.put(mainOutput.getId(), mainOutput);
    }
    for (TupleTag<?> additionalTag : additionalOutputs) {
      outputMap.put(additionalTag.getId(), additionalTag);
    }
    return outputMap;
  }
}