static <T, W extends BoundedWindow>
      ThrowingFunction<KV<T, Iterable<W>>, KV<T, KV<Iterable<W>, Iterable<KV<W, Iterable<W>>>>>>
          createMapFunctionForPTransform(String ptransformId, PTransform ptransform)
              throws IOException {
    RunnerApi.FunctionSpec payload =
        RunnerApi.FunctionSpec.parseFrom(ptransform.getSpec().getPayload());

    WindowFn<?, W> windowFn =
        (WindowFn<?, W>) WindowingStrategyTranslation.windowFnFromProto(payload);
    return WindowMergingFnRunner.<T, W>create(windowFn)::mergeWindows;
  }
@Test
  public void testWindowMergingWithMergingWindowFn() throws Exception {
    ThrowingFunction<
            KV<Object, Iterable<BoundedWindow>>,
            KV<
                Object,
                KV<Iterable<BoundedWindow>, Iterable<KV<BoundedWindow, Iterable<BoundedWindow>>>>>>
        mapFunction =
            WindowMergingFnRunner.createMapFunctionForPTransform(
                "ptransformId",
                createMergeTransformForWindowFn(Sessions.withGapDuration(Duration.millis(5L))));

    // 7, 8 and 10 should all be merged. 1 and 20 should remain in the original set.
    BoundedWindow[] expectedToBeMerged =
        new BoundedWindow[] {
          new IntervalWindow(new Instant(9L), new Instant(11L)),
          new IntervalWindow(new Instant(10L), new Instant(10L)),
          new IntervalWindow(new Instant(7L), new Instant(10L))
        };
    Iterable<BoundedWindow> expectedToBeUnmerged =
        Sets.newHashSet(
            new IntervalWindow(new Instant(1L), new Instant(1L)),
            new IntervalWindow(new Instant(20L), new Instant(20L)));
    KV<Object, Iterable<BoundedWindow>> input =
        KV.of(
            "abc",
            ImmutableList.<BoundedWindow>builder()
                .add(expectedToBeMerged)
                .addAll(expectedToBeUnmerged)
                .build());

    KV<Object, KV<Iterable<BoundedWindow>, Iterable<KV<BoundedWindow, Iterable<BoundedWindow>>>>>
        output = mapFunction.apply(input);
    assertEquals(input.getKey(), output.getKey());
    assertEquals(expectedToBeUnmerged, output.getValue().getKey());
    KV<BoundedWindow, Iterable<BoundedWindow>> mergedOutput =
        Iterables.getOnlyElement(output.getValue().getValue());
    assertEquals(new IntervalWindow(new Instant(7L), new Instant(11L)), mergedOutput.getKey());
    assertThat(mergedOutput.getValue(), containsInAnyOrder(expectedToBeMerged));

    // Process a new group of windows, make sure that previous result has been cleaned up.
    BoundedWindow[] expectedToBeMergedGroup2 =
        new BoundedWindow[] {
          new IntervalWindow(new Instant(15L), new Instant(17L)),
          new IntervalWindow(new Instant(16L), new Instant(18L))
        };

    input =
        KV.of(
            "abc",
            ImmutableList.<BoundedWindow>builder()
                .add(expectedToBeMergedGroup2)
                .addAll(expectedToBeUnmerged)
                .build());

    output = mapFunction.apply(input);
    assertEquals(input.getKey(), output.getKey());
    assertEquals(expectedToBeUnmerged, output.getValue().getKey());
    mergedOutput = Iterables.getOnlyElement(output.getValue().getValue());
    assertEquals(new IntervalWindow(new Instant(15L), new Instant(18L)), mergedOutput.getKey());
    assertThat(mergedOutput.getValue(), containsInAnyOrder(expectedToBeMergedGroup2));
  }