@Override
  public PTransformReplacement<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>
      getReplacementTransform(
          AppliedPTransform<
                  PCollection<KV<K, V>>,
                  PCollection<KV<K, Iterable<V>>>,
                  PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>>
              transform) {

    PCollection<KV<K, Iterable<V>>> output =
        (PCollection<KV<K, Iterable<V>>>) Iterables.getOnlyElement(transform.getOutputs().values());

    return PTransformReplacement.of(
        PTransformReplacements.getSingletonMainInput(transform),
        new DirectGroupByKey<>(transform.getTransform(), output.getWindowingStrategy()));
  }
@Test
  public void getInputSucceeds() {
    TestPipeline p = TestPipeline.create();
    PCollection<KV<String, Integer>> input =
        p.apply(
            Create.of(KV.of("foo", 1))
                .withCoder(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())));
    PCollection<KV<String, Iterable<Integer>>> grouped = input.apply(GroupByKey.create());
    AppliedPTransform<?, ?, ?> producer = DirectGraphs.getProducer(grouped);
    PTransformReplacement<
            PCollection<KV<String, Integer>>, PCollection<KV<String, Iterable<Integer>>>>
        replacement = factory.getReplacementTransform((AppliedPTransform) producer);
    assertThat(replacement.getInput(), Matchers.<PCollection<?>>equalTo(input));
  }