DirectGroupByKeyOverrideFactory extends SingleInputOutputOverrideFactory<
        PCollection<KV<K, V>>,
        PCollection<KV<K, Iterable<V>>>,
        PTransform<PCollection<KV<K, V>>, PCollection<KV<K, Iterable<V>>>>> { @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())); }  @Override 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);  }
@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)); }