/*
 * 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.nemo.compiler.frontend.beam.transform;

import com.google.common.collect.Iterables;
import junit.framework.TestCase;
import org.apache.beam.runners.core.SystemReduceFn;
import org.apache.beam.sdk.coders.*;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.DoFnSchemaInformation;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.windowing.*;
import org.apache.beam.sdk.util.AppliedCombineFn;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.*;
import org.apache.nemo.common.ir.vertex.transform.Transform;
import org.apache.nemo.common.punctuation.Watermark;
import org.apache.nemo.compiler.frontend.beam.NemoPipelineOptions;
import org.joda.time.Duration;
import org.joda.time.Instant;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.util.*;

import static org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing.*;
import static org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode.ACCUMULATING_FIRED_PANES;
import static org.mockito.Mockito.mock;

public class GBKTransformTest extends TestCase {
  private static final Logger LOG = LoggerFactory.getLogger(GBKTransformTest.class.getName());
  private final static Coder STRING_CODER = StringUtf8Coder.of();
  private final static Coder INTEGER_CODER = BigEndianIntegerCoder.of();

  private void checkOutput(final KV<String, Integer> expected, final KV<String, Integer> result) {
    // check key
    assertEquals(expected.getKey(), result.getKey());
    // check value
    assertEquals(expected.getValue(), result.getValue());
  }

  private void checkOutput2(final KV<String, List<String>> expected, final KV<String, Iterable<String>> result) {
    // check key
    assertEquals(expected.getKey(), result.getKey());
    // check value
    final List<String> resultValue = new ArrayList<>();
    final List<String> expectedValue = new ArrayList<>(expected.getValue());
    result.getValue().iterator().forEachRemaining(resultValue::add);
    Collections.sort(resultValue);
    Collections.sort(expectedValue);
    assertEquals(expectedValue, resultValue);
  }


  // Test Combine.Perkey operation.

  // Define combine function.
  public static class CountFn extends Combine.CombineFn<Integer, CountFn.Accum, Integer> {

    public static class Accum {
      int sum = 0;
    }

    @Override
    public Accum createAccumulator() {
      return new Accum();
    }

    @Override
    public Accum addInput(Accum accum, Integer input) {
      accum.sum += input;
      return accum;
    }

    @Override
    public Accum mergeAccumulators(Iterable<Accum> accums) {
      Accum merged = createAccumulator();
      for (Accum accum : accums) {
        merged.sum += accum.sum;
      }
      return merged;
    }

    @Override
    public Integer extractOutput(Accum accum) {
      return accum.sum;
    }

    @Override
    public Coder<CountFn.Accum> getAccumulatorCoder(CoderRegistry registry, Coder<Integer> inputcoder) {
      return AvroCoder.of(CountFn.Accum.class);
    }
  }

  public final static Combine.CombineFn combine_fn = new CountFn();

  // window size: 10 sec
  // period: 5 sec
  //
  // [----------------------- window1 --------------------------]
  //                             [---------------------------window2-------------------------]
  //                                                            [-------------------------window3----------------------]
  //
  //  ts1 -- ts2 -------------------- ts3 -- w1 -- ts4 --------- ts5 - w2 ------------ ts6 -----ts7 -- w3 -- ts8 --ts9 --- w4

  // Test without late data
  @Test
  @SuppressWarnings("unchecked")
  public void test_combine() {
    final TupleTag<String> outputTag = new TupleTag<>("main-output");
    final SlidingWindows slidingWindows = SlidingWindows.of(Duration.standardSeconds(10))
      .every(Duration.standardSeconds(5));

    final Instant ts1 = new Instant(1000);
    final Instant ts2 = new Instant(2000);
    final Instant ts3 = new Instant(6000);
    final Instant ts4 = new Instant(8000);
    final Instant ts5 = new Instant(11000);
    final Instant ts6 = new Instant(14000);
    final Instant ts7 = new Instant(16000);
    final Instant ts8 = new Instant(17000);
    final Instant ts9 = new Instant(19000);
    final Watermark watermark1 = new Watermark(7000);
    final Watermark watermark2 = new Watermark(12000);
    final Watermark watermark3 = new Watermark(18000);
    final Watermark watermark4 = new Watermark(21000);

    AppliedCombineFn<String, Integer, CountFn.Accum, Integer> applied_combine_fn =
      AppliedCombineFn.withInputCoder(
        combine_fn,
        CoderRegistry.createDefault(),
        KvCoder.of(STRING_CODER, INTEGER_CODER),
        null,
        WindowingStrategy.of(slidingWindows).withMode(ACCUMULATING_FIRED_PANES)
      );

    final GBKTransform<String, Integer, Integer> combine_transform =
      new GBKTransform(
        KvCoder.of(STRING_CODER, INTEGER_CODER),
        Collections.singletonMap(outputTag, KvCoder.of(STRING_CODER, INTEGER_CODER)),
        outputTag,
        WindowingStrategy.of(slidingWindows).withMode(ACCUMULATING_FIRED_PANES),
        PipelineOptionsFactory.as(NemoPipelineOptions.class),
        SystemReduceFn.combining(STRING_CODER, applied_combine_fn),
        DoFnSchemaInformation.create(),
        DisplayData.none(),
        false);

    // window1 : [-5000, 5000) in millisecond
    // window2 : [0, 10000)
    // window3 : [5000, 15000)
    // window4 : [10000, 20000)
    List<IntervalWindow> sortedWindows = new ArrayList<>(slidingWindows.assignWindows(ts1));
    Collections.sort(sortedWindows, IntervalWindow::compareTo);
    final IntervalWindow window1 = sortedWindows.get(0);
    final IntervalWindow window2 = sortedWindows.get(1);
    sortedWindows = new ArrayList<>(slidingWindows.assignWindows(ts5));
    Collections.sort(sortedWindows, IntervalWindow::compareTo);
    final IntervalWindow window3 = sortedWindows.get(0);
    final IntervalWindow window4 = sortedWindows.get(1);

    // Prepare to test CombineStreamTransform
    final Transform.Context context = mock(Transform.Context.class);
    final TestOutputCollector<KV<String, Integer>> oc = new TestOutputCollector();
    combine_transform.prepare(context, oc);

    combine_transform.onData(WindowedValue.of(
      KV.of("a", 1), ts1, slidingWindows.assignWindows(ts1), PaneInfo.NO_FIRING));
    combine_transform.onData(WindowedValue.of(
      KV.of("c", 1), ts2, slidingWindows.assignWindows(ts2), PaneInfo.NO_FIRING));
    combine_transform.onData(WindowedValue.of(
      KV.of("b", 1), ts3, slidingWindows.assignWindows(ts3), PaneInfo.NO_FIRING));

    // Emit outputs of window1
    combine_transform.onWatermark(watermark1);
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));

    // Check outputs
    assertEquals(Arrays.asList(window1), oc.outputs.get(0).getWindows());
    assertEquals(2, oc.outputs.size());
    checkOutput(KV.of("a", 1), oc.outputs.get(0).getValue());
    checkOutput(KV.of("c", 1), oc.outputs.get(1).getValue());
    oc.outputs.clear();
    oc.watermarks.clear();

    combine_transform.onData(WindowedValue.of(
      KV.of("a", 1), ts4, slidingWindows.assignWindows(ts4), PaneInfo.NO_FIRING));
    combine_transform.onData(WindowedValue.of(
      KV.of("c", 1), ts5, slidingWindows.assignWindows(ts5), PaneInfo.NO_FIRING));

    // Emit outputs of window2
    combine_transform.onWatermark(watermark2);
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));

    // Check outputs
    assertEquals(Arrays.asList(window2), oc.outputs.get(0).getWindows());
    assertEquals(3, oc.outputs.size());
    checkOutput(KV.of("a", 2), oc.outputs.get(0).getValue());
    checkOutput(KV.of("b", 1), oc.outputs.get(1).getValue());
    checkOutput(KV.of("c", 1), oc.outputs.get(2).getValue());
    oc.outputs.clear();
    oc.watermarks.clear();

    combine_transform.onData(WindowedValue.of(
      KV.of("b", 1), ts6, slidingWindows.assignWindows(ts6), PaneInfo.NO_FIRING));
    combine_transform.onData(WindowedValue.of(
      KV.of("b", 1), ts7, slidingWindows.assignWindows(ts7), PaneInfo.NO_FIRING));
    combine_transform.onData(WindowedValue.of(
      KV.of("a", 1), ts8, slidingWindows.assignWindows(ts8), PaneInfo.NO_FIRING));

    // Emit outputs of window3
    combine_transform.onWatermark(watermark3);
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));

    // Check outputs
    assertEquals(Arrays.asList(window3), oc.outputs.get(0).getWindows());
    checkOutput(KV.of("a", 1), oc.outputs.get(0).getValue());
    checkOutput(KV.of("b", 2), oc.outputs.get(1).getValue());
    checkOutput(KV.of("c", 1), oc.outputs.get(2).getValue());
    oc.outputs.clear();
    oc.watermarks.clear();


    combine_transform.onData(WindowedValue.of(
      KV.of("c", 3), ts9, slidingWindows.assignWindows(ts9), PaneInfo.NO_FIRING));

    // Emit outputs of window3
    combine_transform.onWatermark(watermark4);
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));

    // Check outputs
    assertEquals(Arrays.asList(window4), oc.outputs.get(0).getWindows());
    checkOutput(KV.of("a", 1), oc.outputs.get(0).getValue());
    checkOutput(KV.of("b", 2), oc.outputs.get(1).getValue());
    checkOutput(KV.of("c", 4), oc.outputs.get(2).getValue());

    oc.outputs.clear();
    oc.watermarks.clear();
  }

  // Test with late data
  @Test
  @SuppressWarnings("unchecked")
  public void test_combine_lateData() {
    final TupleTag<String> outputTag = new TupleTag<>("main-output");
    final Duration lateness = Duration.standardSeconds(2);
    final SlidingWindows slidingWindows = SlidingWindows.of(Duration.standardSeconds(10))
      .every(Duration.standardSeconds(5));

    final Instant ts1 = new Instant(1000);
    final Instant ts2 = new Instant(2000);
    final Instant ts3 = new Instant(4500);
    final Instant ts4 = new Instant(11000);
    final Watermark watermark1 = new Watermark(6500);
    final Watermark watermark2 = new Watermark(8000);

    AppliedCombineFn<String, Integer, CountFn.Accum, Integer> applied_combine_fn =
      AppliedCombineFn.withInputCoder(
        combine_fn,
        CoderRegistry.createDefault(),
        KvCoder.of(STRING_CODER, INTEGER_CODER),
        null,
        WindowingStrategy.of(slidingWindows).withMode(ACCUMULATING_FIRED_PANES).withAllowedLateness(lateness)
      );

    final GBKTransform<String, Integer, Integer> combine_transform =
      new GBKTransform(
        KvCoder.of(STRING_CODER, INTEGER_CODER),
        Collections.singletonMap(outputTag, KvCoder.of(STRING_CODER, INTEGER_CODER)),
        outputTag,
        WindowingStrategy.of(slidingWindows).withMode(ACCUMULATING_FIRED_PANES).withAllowedLateness(lateness),
        PipelineOptionsFactory.as(NemoPipelineOptions.class),
        SystemReduceFn.combining(STRING_CODER, applied_combine_fn),
        DoFnSchemaInformation.create(),
        DisplayData.none(),
        false);

    // window1 : [-5000, 5000) in millisecond
    // window2 : [0, 10000)
    // window3 : [5000, 15000)
    // window4 : [10000, 20000)
    List<IntervalWindow> sortedWindows = new ArrayList<>(slidingWindows.assignWindows(ts1));
    Collections.sort(sortedWindows, IntervalWindow::compareTo);
    final IntervalWindow window1 = sortedWindows.get(0);
    final IntervalWindow window2 = sortedWindows.get(1);
    sortedWindows = new ArrayList<>(slidingWindows.assignWindows(ts4));
    Collections.sort(sortedWindows, IntervalWindow::compareTo);
    final IntervalWindow window3 = sortedWindows.get(0);
    final IntervalWindow window4 = sortedWindows.get(1);

    // Prepare to test
    final Transform.Context context = mock(Transform.Context.class);
    final TestOutputCollector<KV<String, Integer>> oc = new TestOutputCollector();
    combine_transform.prepare(context, oc);

    combine_transform.onData(WindowedValue.of(
      KV.of("a", 1), ts1, slidingWindows.assignWindows(ts1), PaneInfo.NO_FIRING));
    combine_transform.onData(WindowedValue.of(
      KV.of("b", 1), ts2, slidingWindows.assignWindows(ts2), PaneInfo.NO_FIRING));

    // On-time firing of window1. Skipping checking outputs since test1 checks output from non-late data
    combine_transform.onWatermark(watermark1);
    oc.outputs.clear();

    // Late data in window 1. Should be accumulated since EOW + allowed lateness > current Watermark
    combine_transform.onData(WindowedValue.of(
      KV.of("a", 5), ts1, slidingWindows.assignWindows(ts1), PaneInfo.NO_FIRING));

    // Check outputs
    assertEquals(Arrays.asList(window1), oc.outputs.get(0).getWindows());
    assertEquals(1,oc.outputs.size());
    assertEquals(LATE, oc.outputs.get(0).getPane().getTiming());
    checkOutput(KV.of("a", 6), oc.outputs.get(0).getValue());

    oc.outputs.clear();
    oc.watermarks.clear();

    // Late data in window 1. Should NOT be accumulated to outputs of window1 since EOW + allowed lateness > current Watermark
    combine_transform.onWatermark(watermark2);
    combine_transform.onData(WindowedValue.of(
      KV.of("a", 10), ts3, slidingWindows.assignWindows(ts3), PaneInfo.NO_FIRING));
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));


    // Check outputs
    assertEquals(Arrays.asList(window1), oc.outputs.get(0).getWindows());
    assertEquals(1, oc.outputs.size());
    assertEquals(LATE, oc.outputs.get(0).getPane().getTiming());
    checkOutput(KV.of("a", 10), oc.outputs.get(0).getValue());
    oc.outputs.clear();
    oc.watermarks.clear();
  }

  // Now testing GroupbyKey Operation.

  // window size: 2 sec
  // interval size: 1 sec
  //
  //                           [--------------window2------------------------------]
  // [----------------------- window1 --------------------------]
  // [-------window0-------]
  // ts1 -- ts2 -- ts3 -- w -- ts4 -- w2 -- ts5 --ts6 --ts7 -- w3 -- ts8 --ts9 - --w4
  // (1, "hello")
  //      (1, "world")
  //             (2, "hello")
  //                   ==> window1: {(1,["hello","world"]), (2, ["hello"])}
  //                                 (1, "a")
  //                                                       (2,"a")
  //                                                             (3,"a")
  //                                                                  (2,"b")
  //                                                       => window2: {(1,"a"), (2,["a","b"]), (3,"a")}

  @Test
  @SuppressWarnings("unchecked")
  public void test_gbk() {

    final TupleTag<String> outputTag = new TupleTag<>("main-output");
    final SlidingWindows slidingWindows = SlidingWindows.of(Duration.standardSeconds(2))
      .every(Duration.standardSeconds(1));

    final GBKTransform<String, String, Iterable<String>> doFnTransform =
      new GBKTransform(
        KvCoder.of(STRING_CODER, STRING_CODER),
        Collections.singletonMap(outputTag, KvCoder.of(STRING_CODER, IterableCoder.of(STRING_CODER))),
        outputTag,
        WindowingStrategy.of(slidingWindows),
        PipelineOptionsFactory.as(NemoPipelineOptions.class),
        SystemReduceFn.buffering(STRING_CODER),
        DoFnSchemaInformation.create(),
        DisplayData.none(),
        false);

    final Instant ts1 = new Instant(1);
    final Instant ts2 = new Instant(100);
    final Instant ts3 = new Instant(300);
    final Watermark watermark = new Watermark(1003);
    final Instant ts4 = new Instant(1200);
    final Watermark watermark2 = new Watermark(1400);
    final Instant ts5 = new Instant(1600);
    final Instant ts6 = new Instant(1800);
    final Instant ts7 = new Instant(1900);
    final Watermark watermark3 = new Watermark(2100);
    final Instant ts8 = new Instant(2200);
    final Instant ts9 = new Instant(2300);
    final Watermark watermark4 = new Watermark(3000);


    List<IntervalWindow> sortedWindows = new ArrayList<>(slidingWindows.assignWindows(ts1));
    Collections.sort(sortedWindows, IntervalWindow::compareTo);

    // [0---1000)
    final IntervalWindow window0 = sortedWindows.get(0);
    // [0---2000)
    final IntervalWindow window1 = sortedWindows.get(1);

    sortedWindows.clear();
    sortedWindows = new ArrayList<>(slidingWindows.assignWindows(ts4));
    Collections.sort(sortedWindows, IntervalWindow::compareTo);

    // [1000--3000)
    final IntervalWindow window2 = sortedWindows.get(1);


    final Transform.Context context = mock(Transform.Context.class);
    final TestOutputCollector<KV<String, Iterable<String>>> oc = new TestOutputCollector();
    doFnTransform.prepare(context, oc);

    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "hello"), ts1, slidingWindows.assignWindows(ts1), PaneInfo.NO_FIRING));

    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "world"), ts2, slidingWindows.assignWindows(ts2), PaneInfo.NO_FIRING));

    doFnTransform.onData(WindowedValue.of(
      KV.of("2", "hello"), ts3, slidingWindows.assignWindows(ts3), PaneInfo.NO_FIRING));

    doFnTransform.onWatermark(watermark);

    // output
    // 1: ["hello", "world"]
    // 2: ["hello"]
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));

    // windowed result for key 1
    assertEquals(Arrays.asList(window0), oc.outputs.get(0).getWindows());
    checkOutput2(KV.of("1", Arrays.asList("hello", "world")), oc.outputs.get(0).getValue());

    // windowed result for key 2
    assertEquals(Arrays.asList(window0), oc.outputs.get(1).getWindows());
    checkOutput2(KV.of("2", Arrays.asList("hello")), oc.outputs.get(1).getValue());

    assertEquals(2, oc.outputs.size());
    assertEquals(2, oc.watermarks.size());

    // check output watermark
    assertEquals(1000,
      oc.watermarks.get(0).getTimestamp());

    oc.outputs.clear();
    oc.watermarks.clear();

    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "a"), ts4, slidingWindows.assignWindows(ts4), PaneInfo.NO_FIRING));


    doFnTransform.onWatermark(watermark2);

    assertEquals(0, oc.outputs.size()); // do not emit anything
    assertEquals(0, oc.watermarks.size());

    doFnTransform.onData(WindowedValue.of(
      KV.of("3", "a"), ts5, slidingWindows.assignWindows(ts5), PaneInfo.NO_FIRING));

    doFnTransform.onData(WindowedValue.of(
      KV.of("3", "a"), ts6, slidingWindows.assignWindows(ts6), PaneInfo.NO_FIRING));

    doFnTransform.onData(WindowedValue.of(
      KV.of("3", "b"), ts7, slidingWindows.assignWindows(ts7), PaneInfo.NO_FIRING));

    // emit window1
    doFnTransform.onWatermark(watermark3);

    // output
    // 1: ["hello", "world", "a"]
    // 2: ["hello"]
    // 3: ["a", "a", "b"]
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));


    // windowed result for key 1
    assertEquals(Arrays.asList(window1), oc.outputs.get(0).getWindows());
    checkOutput2(KV.of("1", Arrays.asList("hello", "world", "a")), oc.outputs.get(0).getValue());

    // windowed result for key 2
    assertEquals(Arrays.asList(window1), oc.outputs.get(1).getWindows());
    checkOutput2(KV.of("2", Arrays.asList("hello")), oc.outputs.get(1).getValue());

    // windowed result for key 3
    assertEquals(Arrays.asList(window1), oc.outputs.get(2).getWindows());
    checkOutput2(KV.of("3", Arrays.asList("a", "a", "b")), oc.outputs.get(2).getValue());

    // check output watermark
    assertEquals(2000,
      oc.watermarks.get(0).getTimestamp());

    oc.outputs.clear();
    oc.watermarks.clear();


    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "a"), ts8, slidingWindows.assignWindows(ts8), PaneInfo.NO_FIRING));

    doFnTransform.onData(WindowedValue.of(
      KV.of("3", "b"), ts9, slidingWindows.assignWindows(ts9), PaneInfo.NO_FIRING));

    // emit window2
    doFnTransform.onWatermark(watermark4);

    // output
    // 1: ["a", "a"]
    // 3: ["a", "a", "b", "b"]
    Collections.sort(oc.outputs, (o1, o2) -> o1.getValue().getKey().compareTo(o2.getValue().getKey()));

    assertEquals(2, oc.outputs.size());

    // windowed result for key 1
    assertEquals(Arrays.asList(window2), oc.outputs.get(0).getWindows());
    checkOutput2(KV.of("1", Arrays.asList("a", "a")), oc.outputs.get(0).getValue());

    // windowed result for key 3
    assertEquals(Arrays.asList(window2), oc.outputs.get(1).getWindows());
    checkOutput2(KV.of("3", Arrays.asList("a", "a", "b", "b")), oc.outputs.get(1).getValue());

    // check output watermark
    assertEquals(3000,
      oc.watermarks.get(0).getTimestamp());

    doFnTransform.close();
  }

  /**
   * Test complex triggers that emit early and late firing.
   */
  @Test
  public void test_gbk_eventTimeTrigger() {
    final Duration lateness = Duration.standardSeconds(1);
    final AfterWatermark.AfterWatermarkEarlyAndLate trigger = AfterWatermark.pastEndOfWindow()
      // early firing
      .withEarlyFirings(
        AfterProcessingTime
          .pastFirstElementInPane()
          // early firing 1 sec after receiving an element
          .plusDelayOf(Duration.millis(1000)))
      // late firing: Fire on any late data.
      .withLateFirings(AfterPane.elementCountAtLeast(1));

    final FixedWindows window = (FixedWindows) Window.into(
      FixedWindows.of(Duration.standardSeconds(5)))
      // lateness
      .withAllowedLateness(lateness)
      .triggering(trigger)
      // TODO #308: Test discarding of refinements
      .accumulatingFiredPanes().getWindowFn();

    final TupleTag<String> outputTag = new TupleTag<>("main-output");

    final GBKTransform<String, String, Iterable<String>> doFnTransform =
      new GBKTransform(
        KvCoder.of(STRING_CODER, STRING_CODER),
        Collections.singletonMap(outputTag, KvCoder.of(STRING_CODER, IterableCoder.of(STRING_CODER))),
        outputTag,
        WindowingStrategy.of(window).withTrigger(trigger)
          .withMode(ACCUMULATING_FIRED_PANES)
          .withAllowedLateness(lateness),
        PipelineOptionsFactory.as(NemoPipelineOptions.class),
        SystemReduceFn.buffering(STRING_CODER),
        DoFnSchemaInformation.create(),
        DisplayData.none(),
        false);


    final Transform.Context context = mock(Transform.Context.class);
    final TestOutputCollector<KV<String, Iterable<String>>> oc = new TestOutputCollector();
    doFnTransform.prepare(context, oc);

    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "hello"), new Instant(1), window.assignWindow(new Instant(1)), PaneInfo.NO_FIRING));

    try {
      Thread.sleep(1000);
    } catch (InterruptedException e) {
      e.printStackTrace();
    }

    // early firing is not related to the watermark progress
    doFnTransform.onWatermark(new Watermark(2));
    assertEquals(1, oc.outputs.size());
    assertEquals(EARLY, oc.outputs.get(0).getPane().getTiming());
    oc.outputs.clear();

    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "world"), new Instant(3), window.assignWindow(new Instant(3)), PaneInfo.NO_FIRING));
    // EARLY firing... waiting >= 1 sec
    try {
      Thread.sleep(2000);
    } catch (InterruptedException e) {
      e.printStackTrace();
    }

    // GBKTransform emits data when receiving watermark
    // TODO #250: element-wise processing
    doFnTransform.onWatermark(new Watermark(5));
    assertEquals(1, oc.outputs.size());
    assertEquals(EARLY, oc.outputs.get(0).getPane().getTiming());
    // ACCUMULATION MODE
    checkOutput2(KV.of("1", Arrays.asList("hello", "world")), oc.outputs.get(0).getValue());
    oc.outputs.clear();

    // ON TIME
    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "!!"), new Instant(3), window.assignWindow(new Instant(3)), PaneInfo.NO_FIRING));
    doFnTransform.onWatermark(new Watermark(5001));
    assertEquals(1, oc.outputs.size());
    assertEquals(ON_TIME, oc.outputs.get(0).getPane().getTiming());
    // ACCUMULATION MODE
    checkOutput2(KV.of("1", Arrays.asList("hello", "world", "!!")), oc.outputs.get(0).getValue());
    oc.outputs.clear();

    // LATE DATA
    // actual window: [0-5000)
    // allowed lateness: 1000 (ms)
    // current watermark: 5001
    // data: 1000
    // End of current window + allowed lateness >= current watermark (4999 + 1000 >= 5001)
    // so it should be accumulated to the prev window
    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "bye!"), new Instant(1000),
      window.assignWindow(new Instant(1000)), PaneInfo.NO_FIRING));
    doFnTransform.onWatermark(new Watermark(6000));
    assertEquals(1, oc.outputs.size());
    assertEquals(LATE, oc.outputs.get(0).getPane().getTiming());
    // The data should  be accumulated to the previous window because it allows 1 second lateness
    checkOutput2(KV.of("1", Arrays.asList("hello", "world", "!!", "bye!")), oc.outputs.get(0).getValue());
    oc.outputs.clear();

    // LATE DATA
    // actual window: [0-5000)
    // allowed lateness: 1000 (ms)
    // data timestamp: 4800
    // current watermark: 6000
    // End of current window + allowed lateness < current watermark (4999 + 1000 < 6000)
    // It should not be accumulated to the prev window
    doFnTransform.onData(WindowedValue.of(
      KV.of("1", "hello again!"), new Instant(4800),
      window.assignWindow(new Instant(4800)), PaneInfo.NO_FIRING));
    doFnTransform.onWatermark(new Watermark(6300));
    assertEquals(1, oc.outputs.size());
    assertEquals(LATE, oc.outputs.get(0).getPane().getTiming());
    checkOutput2(KV.of("1", Arrays.asList("hello again!")), oc.outputs.get(0).getValue());
    oc.outputs.clear();
    doFnTransform.close();
  }
}