@Override
  public <InputT> TransformEvaluator<InputT> forApplication(
      AppliedPTransform<?, ?, ?> application, CommittedBundle<?> inputBundle) {
    @SuppressWarnings({"cast", "unchecked", "rawtypes"})
    TransformEvaluator<InputT> evaluator = createEvaluator((AppliedPTransform) application);
    return evaluator;
  }
@Test
  public void testInMemoryEvaluator() throws Exception {
    KV<String, Integer> firstFoo = KV.of("foo", -1);
    KV<String, Integer> secondFoo = KV.of("foo", 1);
    KV<String, Integer> thirdFoo = KV.of("foo", 3);
    KV<String, Integer> firstBar = KV.of("bar", 22);
    KV<String, Integer> secondBar = KV.of("bar", 12);
    KV<String, Integer> firstBaz = KV.of("baz", Integer.MAX_VALUE);
    PCollection<KV<String, Integer>> values =
        p.apply(Create.of(firstFoo, firstBar, secondFoo, firstBaz, secondBar, thirdFoo));
    PCollection<KeyedWorkItem<String, Integer>> groupedKvs =
        values.apply(new DirectGroupByKeyOnly<>());

    CommittedBundle<KV<String, Integer>> inputBundle =
        bundleFactory.createBundle(values).commit(Instant.now());
    EvaluationContext evaluationContext = mock(EvaluationContext.class);

    StructuralKey<String> fooKey = StructuralKey.of("foo", StringUtf8Coder.of());
    UncommittedBundle<KeyedWorkItem<String, Integer>> fooBundle =
        bundleFactory.createKeyedBundle(fooKey, groupedKvs);
    StructuralKey<String> barKey = StructuralKey.of("bar", StringUtf8Coder.of());
    UncommittedBundle<KeyedWorkItem<String, Integer>> barBundle =
        bundleFactory.createKeyedBundle(barKey, groupedKvs);
    StructuralKey<String> bazKey = StructuralKey.of("baz", StringUtf8Coder.of());
    UncommittedBundle<KeyedWorkItem<String, Integer>> bazBundle =
        bundleFactory.createKeyedBundle(bazKey, groupedKvs);

    when(evaluationContext.createKeyedBundle(fooKey, groupedKvs)).thenReturn(fooBundle);
    when(evaluationContext.createKeyedBundle(barKey, groupedKvs)).thenReturn(barBundle);
    when(evaluationContext.createKeyedBundle(bazKey, groupedKvs)).thenReturn(bazBundle);

    // The input to a GroupByKey is assumed to be a KvCoder
    @SuppressWarnings("unchecked")
    Coder<String> keyCoder = ((KvCoder<String, Integer>) values.getCoder()).getKeyCoder();
    TransformEvaluator<KV<String, Integer>> evaluator =
        new GroupByKeyOnlyEvaluatorFactory(evaluationContext)
            .forApplication(DirectGraphs.getProducer(groupedKvs), inputBundle);

    evaluator.processElement(WindowedValue.valueInGlobalWindow(firstFoo));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(secondFoo));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(thirdFoo));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(firstBar));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(secondBar));
    evaluator.processElement(WindowedValue.valueInGlobalWindow(firstBaz));

    evaluator.finishBundle();

    assertThat(
        fooBundle.commit(Instant.now()).getElements(),
        contains(
            new KeyedWorkItemMatcher<>(
                KeyedWorkItems.elementsWorkItem(
                    "foo",
                    ImmutableSet.of(
                        WindowedValue.valueInGlobalWindow(-1),
                        WindowedValue.valueInGlobalWindow(1),
                        WindowedValue.valueInGlobalWindow(3))),
                keyCoder)));
    assertThat(
        barBundle.commit(Instant.now()).getElements(),
        contains(
            new KeyedWorkItemMatcher<>(
                KeyedWorkItems.elementsWorkItem(
                    "bar",
                    ImmutableSet.of(
                        WindowedValue.valueInGlobalWindow(12),
                        WindowedValue.valueInGlobalWindow(22))),
                keyCoder)));
    assertThat(
        bazBundle.commit(Instant.now()).getElements(),
        contains(
            new KeyedWorkItemMatcher<>(
                KeyedWorkItems.elementsWorkItem(
                    "baz", ImmutableSet.of(WindowedValue.valueInGlobalWindow(Integer.MAX_VALUE))),
                keyCoder)));
  }