/*
 * 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.dataflow.worker;

import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.equalTo;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Matchers.eq;
import static org.mockito.Mockito.when;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.beam.runners.core.DoFnRunner;
import org.apache.beam.runners.core.DoFnRunners;
import org.apache.beam.runners.core.InMemoryStateInternals;
import org.apache.beam.runners.core.SideInputReader;
import org.apache.beam.runners.core.StateInternals;
import org.apache.beam.runners.core.StateNamespaces;
import org.apache.beam.runners.dataflow.worker.StateFetcher.SideInputState;
import org.apache.beam.runners.dataflow.worker.util.ListOutputManager;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill;
import org.apache.beam.runners.dataflow.worker.windmill.Windmill.GlobalDataRequest;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.testing.TestPipeline;
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.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.IntervalWindow;
import org.apache.beam.sdk.transforms.windowing.PaneInfo;
import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollectionView;
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.guava.v26_0_jre.com.google.common.collect.Iterables;
import org.hamcrest.Matchers;
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.Mock;
import org.mockito.Mockito;
import org.mockito.MockitoAnnotations;

/** Unit tests for {@link StreamingSideInputDoFnRunner}. */
@RunWith(JUnit4.class)
public class StreamingSideInputDoFnRunnerTest {

  private static final FixedWindows WINDOW_FN = FixedWindows.of(Duration.millis(10));

  static TupleTag<String> mainOutputTag = new TupleTag<>();
  @Mock StreamingModeExecutionContext execContext;
  @Mock StreamingModeExecutionContext.StepContext stepContext;
  @Mock SideInputReader mockSideInputReader;

  private final InMemoryStateInternals<String> state = InMemoryStateInternals.forKey("dummyKey");

  // Mockito causes the covariant ? in StepContext#stateInternals to become contravariant
  // because it is accepted as a parameter to #thenReturn. It is completely safe to
  // treat InMemoryStateInternals<String> as StateInternals<?>.
  @SuppressWarnings({"unchecked", "rawtypes"})
  @Before
  public void setUp() {
    MockitoAnnotations.initMocks(this);
    when(stepContext.stateInternals()).thenReturn((StateInternals) state);
  }

  @Test
  public void testSideInputReady() throws Exception {
    PCollectionView<String> view = createView();

    when(stepContext.getSideInputNotifications())
        .thenReturn(Arrays.<Windmill.GlobalDataId>asList());
    when(stepContext.issueSideInputFetch(
            eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN)))
        .thenReturn(true);
    when(execContext.getSideInputReaderForViews(
            Mockito.<Iterable<? extends PCollectionView<?>>>any()))
        .thenReturn(mockSideInputReader);
    when(mockSideInputReader.contains(eq(view))).thenReturn(true);
    when(mockSideInputReader.get(eq(view), any(BoundedWindow.class))).thenReturn("data");

    ListOutputManager outputManager = new ListOutputManager();
    List<PCollectionView<String>> views = Arrays.asList(view);
    StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
    StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner =
        createRunner(outputManager, views, sideInputFetcher);

    runner.startBundle();
    runner.processElement(createDatum("e", 0));
    runner.finishBundle();

    assertThat(outputManager.getOutput(mainOutputTag), contains(createDatum("e:data", 0)));
  }

  @Test
  public void testSideInputNotReady() throws Exception {
    PCollectionView<String> view = createView();

    when(stepContext.getSideInputNotifications())
        .thenReturn(Arrays.<Windmill.GlobalDataId>asList());
    when(stepContext.issueSideInputFetch(
            eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN)))
        .thenReturn(false);

    ListOutputManager outputManager = new ListOutputManager();

    List<PCollectionView<String>> views = Arrays.asList(view);
    StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
    StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner =
        createRunner(outputManager, views, sideInputFetcher);

    runner.startBundle();
    runner.processElement(createDatum("e", 0));
    runner.finishBundle();

    assertTrue(outputManager.getOutput(mainOutputTag).isEmpty());

    IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));

    // Verify that we added the element to an appropriate tag list, and that we buffered the element
    ValueState<Map<IntervalWindow, Set<GlobalDataRequest>>> blockedMapState =
        state.state(
            StateNamespaces.global(),
            StreamingSideInputFetcher.blockedMapAddr(WINDOW_FN.windowCoder()));
    assertEquals(
        blockedMapState.read(),
        Collections.singletonMap(
            window,
            Collections.singleton(
                Windmill.GlobalDataRequest.newBuilder()
                    .setDataId(
                        Windmill.GlobalDataId.newBuilder()
                            .setTag(view.getTagInternal().getId())
                            .setVersion(
                                ByteString.copyFrom(
                                    CoderUtils.encodeToByteArray(
                                        IntervalWindow.getCoder(), window)))
                            .build())
                    .setExistenceWatermarkDeadline(9000)
                    .build())));
    assertThat(
        sideInputFetcher.elementBag(createWindow(0)).read(),
        Matchers.contains(createDatum("e", 0)));
    assertEquals(sideInputFetcher.watermarkHold(createWindow(0)).read(), new Instant(0));
  }

  @Test
  public void testMultipleWindowsNotReady() throws Exception {
    PCollectionView<String> view = createView();

    when(stepContext.getSideInputNotifications())
        .thenReturn(Arrays.<Windmill.GlobalDataId>asList());
    when(stepContext.issueSideInputFetch(
            eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN)))
        .thenReturn(false);

    ListOutputManager outputManager = new ListOutputManager();

    List<PCollectionView<String>> views = Arrays.asList(view);
    StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
    StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner =
        createRunner(
            SlidingWindows.of(Duration.millis(10)).every(Duration.millis(10)),
            outputManager,
            views,
            sideInputFetcher);

    IntervalWindow window1 = new IntervalWindow(new Instant(0), new Instant(10));
    IntervalWindow window2 = new IntervalWindow(new Instant(-5), new Instant(5));
    long timestamp = 1L;

    WindowedValue<String> elem =
        WindowedValue.of(
            "e", new Instant(timestamp), Arrays.asList(window1, window2), PaneInfo.NO_FIRING);

    runner.startBundle();
    runner.processElement(elem);
    runner.finishBundle();

    assertTrue(outputManager.getOutput(mainOutputTag).isEmpty());

    // Verify that we added the element to an appropriate tag list, and that we buffered the element
    // in both windows separately
    ValueState<Map<IntervalWindow, Set<GlobalDataRequest>>> blockedMapState =
        state.state(
            StateNamespaces.global(),
            StreamingSideInputFetcher.blockedMapAddr(WINDOW_FN.windowCoder()));

    Map<IntervalWindow, Set<GlobalDataRequest>> blockedMap = blockedMapState.read();

    assertThat(
        blockedMap.get(window1),
        equalTo(
            Collections.singleton(
                Windmill.GlobalDataRequest.newBuilder()
                    .setDataId(
                        Windmill.GlobalDataId.newBuilder()
                            .setTag(view.getTagInternal().getId())
                            .setVersion(
                                ByteString.copyFrom(
                                    CoderUtils.encodeToByteArray(
                                        IntervalWindow.getCoder(), window1)))
                            .build())
                    .setExistenceWatermarkDeadline(9000)
                    .build())));

    assertThat(
        blockedMap.get(window2),
        equalTo(
            Collections.singleton(
                Windmill.GlobalDataRequest.newBuilder()
                    .setDataId(
                        Windmill.GlobalDataId.newBuilder()
                            .setTag(view.getTagInternal().getId())
                            .setVersion(
                                ByteString.copyFrom(
                                    CoderUtils.encodeToByteArray(
                                        IntervalWindow.getCoder(), window1)))
                            .build())
                    .setExistenceWatermarkDeadline(9000)
                    .build())));

    assertThat(
        sideInputFetcher.elementBag(window1).read(),
        contains(Iterables.get(elem.explodeWindows(), 0)));

    assertThat(
        sideInputFetcher.elementBag(window2).read(),
        contains(Iterables.get(elem.explodeWindows(), 1)));

    assertEquals(sideInputFetcher.watermarkHold(window1).read(), new Instant(timestamp));
    assertEquals(sideInputFetcher.watermarkHold(window2).read(), new Instant(timestamp));
  }

  @Test
  public void testSideInputNotification() throws Exception {
    PCollectionView<String> view = createView();

    IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));
    Windmill.GlobalDataId id =
        Windmill.GlobalDataId.newBuilder()
            .setTag(view.getTagInternal().getId())
            .setVersion(
                ByteString.copyFrom(
                    CoderUtils.encodeToByteArray(IntervalWindow.getCoder(), window)))
            .build();

    Set<Windmill.GlobalDataRequest> requestSet = new HashSet<>();
    requestSet.add(Windmill.GlobalDataRequest.newBuilder().setDataId(id).build());
    Map<IntervalWindow, Set<Windmill.GlobalDataRequest>> blockedMap = new HashMap<>();
    blockedMap.put(window, requestSet);

    ValueState<Map<IntervalWindow, Set<GlobalDataRequest>>> blockedMapState =
        state.state(
            StateNamespaces.global(),
            StreamingSideInputFetcher.blockedMapAddr(WINDOW_FN.windowCoder()));
    blockedMapState.write(blockedMap);

    ListOutputManager outputManager = new ListOutputManager();
    List<PCollectionView<String>> views = Arrays.asList(view);
    StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
    StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner =
        createRunner(outputManager, views, sideInputFetcher);
    sideInputFetcher.watermarkHold(createWindow(0)).add(new Instant(0));
    sideInputFetcher.elementBag(createWindow(0)).add(createDatum("e", 0));

    when(stepContext.getSideInputNotifications()).thenReturn(Arrays.asList(id));
    when(stepContext.issueSideInputFetch(
            eq(view), any(BoundedWindow.class), eq(SideInputState.UNKNOWN)))
        .thenReturn(false);
    when(stepContext.issueSideInputFetch(
            eq(view), any(BoundedWindow.class), eq(SideInputState.KNOWN_READY)))
        .thenReturn(true);
    when(execContext.getSideInputReaderForViews(
            Mockito.<Iterable<? extends PCollectionView<?>>>any()))
        .thenReturn(mockSideInputReader);
    when(mockSideInputReader.contains(eq(view))).thenReturn(true);
    when(mockSideInputReader.get(eq(view), any(BoundedWindow.class))).thenReturn("data");

    runner.startBundle();
    runner.finishBundle();

    assertThat(outputManager.getOutput(mainOutputTag), contains(createDatum("e:data", 0)));

    assertThat(blockedMapState.read(), Matchers.nullValue());
    assertThat(sideInputFetcher.watermarkHold(createWindow(0)).read(), Matchers.nullValue());
    assertThat(sideInputFetcher.elementBag(createWindow(0)).read(), Matchers.emptyIterable());
  }

  @Test
  public void testMultipleSideInputs() throws Exception {
    PCollectionView<String> view1 = createView();
    PCollectionView<String> view2 = createView();

    IntervalWindow window = new IntervalWindow(new Instant(0), new Instant(10));
    Windmill.GlobalDataId id =
        Windmill.GlobalDataId.newBuilder()
            .setTag(view1.getTagInternal().getId())
            .setVersion(
                ByteString.copyFrom(
                    CoderUtils.encodeToByteArray(IntervalWindow.getCoder(), window)))
            .build();

    Set<Windmill.GlobalDataRequest> requestSet = new HashSet<>();
    requestSet.add(Windmill.GlobalDataRequest.newBuilder().setDataId(id).build());
    Map<IntervalWindow, Set<Windmill.GlobalDataRequest>> blockedMap = new HashMap<>();
    blockedMap.put(window, requestSet);

    ValueState<Map<IntervalWindow, Set<GlobalDataRequest>>> blockedMapState =
        state.state(
            StateNamespaces.global(),
            StreamingSideInputFetcher.blockedMapAddr(WINDOW_FN.windowCoder()));
    blockedMapState.write(blockedMap);

    when(stepContext.getSideInputNotifications()).thenReturn(Arrays.asList(id));
    when(stepContext.issueSideInputFetch(
            any(PCollectionView.class), any(BoundedWindow.class), any(SideInputState.class)))
        .thenReturn(true);
    when(execContext.getSideInputReaderForViews(
            Mockito.<Iterable<? extends PCollectionView<?>>>any()))
        .thenReturn(mockSideInputReader);
    when(mockSideInputReader.contains(eq(view1))).thenReturn(true);
    when(mockSideInputReader.contains(eq(view2))).thenReturn(true);
    when(mockSideInputReader.get(eq(view1), any(BoundedWindow.class))).thenReturn("data1");
    when(mockSideInputReader.get(eq(view2), any(BoundedWindow.class))).thenReturn("data2");

    ListOutputManager outputManager = new ListOutputManager();
    List<PCollectionView<String>> views = Arrays.asList(view1, view2);
    StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher = createFetcher(views);
    StreamingSideInputDoFnRunner<String, String, IntervalWindow> runner =
        createRunner(outputManager, views, sideInputFetcher);
    sideInputFetcher.watermarkHold(createWindow(0)).add(new Instant(0));
    sideInputFetcher.elementBag(createWindow(0)).add(createDatum("e1", 0));

    runner.startBundle();
    runner.processElement(createDatum("e2", 2));
    runner.finishBundle();

    assertThat(
        outputManager.getOutput(mainOutputTag),
        contains(createDatum("e1:data1:data2", 0), createDatum("e2:data1:data2", 2)));

    assertThat(blockedMapState.read(), Matchers.nullValue());
    assertThat(sideInputFetcher.watermarkHold(createWindow(0)).read(), Matchers.nullValue());
    assertThat(sideInputFetcher.elementBag(createWindow(0)).read(), Matchers.emptyIterable());
  }

  @SuppressWarnings("unchecked")
  private <ReceiverT> StreamingSideInputDoFnRunner<String, String, IntervalWindow> createRunner(
      DoFnRunners.OutputManager outputManager,
      List<PCollectionView<String>> views,
      StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher)
      throws Exception {
    return createRunner(WINDOW_FN, outputManager, views, sideInputFetcher);
  }

  @SuppressWarnings("unchecked")
  private <ReceiverT> StreamingSideInputDoFnRunner<String, String, IntervalWindow> createRunner(
      WindowFn<?, ?> windowFn,
      DoFnRunners.OutputManager outputManager,
      List<PCollectionView<String>> views,
      StreamingSideInputFetcher<String, IntervalWindow> sideInputFetcher)
      throws Exception {
    DoFnRunner<String, String> simpleDoFnRunner =
        DoFnRunners.simpleRunner(
            PipelineOptionsFactory.create(),
            new SideInputFn(views),
            mockSideInputReader,
            outputManager,
            mainOutputTag,
            Arrays.<TupleTag<?>>asList(),
            stepContext,
            null,
            Collections.emptyMap(),
            WindowingStrategy.of(windowFn),
            DoFnSchemaInformation.create(),
            Collections.emptyMap());
    return new StreamingSideInputDoFnRunner<>(simpleDoFnRunner, sideInputFetcher);
  }

  private <ReceiverT> StreamingSideInputFetcher<String, IntervalWindow> createFetcher(
      List<PCollectionView<String>> views) throws Exception {
    @SuppressWarnings({"unchecked", "rawtypes"})
    Iterable<PCollectionView<?>> typedViews = (Iterable) views;

    return new StreamingSideInputFetcher<>(
        typedViews, StringUtf8Coder.of(), WindowingStrategy.of(WINDOW_FN), stepContext);
  }

  private static class SideInputFn extends DoFn<String, String> {
    private List<PCollectionView<String>> views;

    public SideInputFn(List<PCollectionView<String>> views) {
      this.views = views;
    }

    @ProcessElement
    public void processElement(ProcessContext c) {
      String output = c.element();
      for (PCollectionView<String> view : views) {
        output += ":" + c.sideInput(view);
      }
      c.output(output);
    }
  }

  private PCollectionView<String> createView() {
    return TestPipeline.create()
        .apply(Create.empty(StringUtf8Coder.of()))
        .apply(Window.<String>into(WINDOW_FN))
        .apply(View.<String>asSingleton());
  }

  private WindowedValue<String> createDatum(String element, long timestamp) {
    return WindowedValue.of(
        element,
        new Instant(timestamp),
        Arrays.asList(createWindow(timestamp)),
        PaneInfo.NO_FIRING);
  }

  private IntervalWindow createWindow(long timestamp) {
    return new IntervalWindow(
        new Instant(timestamp - timestamp % 10), new Instant(timestamp - timestamp % 10 + 10));
  }
}