public Set<PValue> getKeyedPValues() {
    checkState(
        finalized, "can't call getKeyedPValues before a Pipeline has been completely traversed");
    return keyedValues;
  }
@Test
  public void keyedInputWithKeyPreserving() {

    PCollection<KV<String, WindowedValue<KV<String, Integer>>>> input =
        p.apply(
            Create.of(
                    KV.of(
                        "hello",
                        WindowedValue.of(
                            KV.of("hello", 3),
                            new Instant(0),
                            new IntervalWindow(new Instant(0), new Instant(9)),
                            PaneInfo.NO_FIRING)))
                .withCoder(
                    KvCoder.of(
                        StringUtf8Coder.of(),
                        WindowedValue.getValueOnlyCoder(
                            KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())))));

    TupleTag<KeyedWorkItem<String, KV<String, Integer>>> keyedTag = new TupleTag<>();
    PCollection<KeyedWorkItem<String, KV<String, Integer>>> keyed =
        input
            .apply(new DirectGroupByKeyOnly<>())
            .apply(
                new DirectGroupAlsoByWindow<>(
                    WindowingStrategy.globalDefault(), WindowingStrategy.globalDefault()))
            .apply(
                ParDo.of(new ParDoMultiOverrideFactory.ToKeyedWorkItem<String, Integer>())
                    .withOutputTags(keyedTag, TupleTagList.empty()))
            .get(keyedTag)
            .setCoder(
                KeyedWorkItemCoder.of(
                    StringUtf8Coder.of(),
                    KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of()),
                    GlobalWindow.Coder.INSTANCE));

    p.traverseTopologically(visitor);
    assertThat(visitor.getKeyedPValues(), hasItem(keyed));
  }