SingleInputOutputOverrideFactory implements PTransformOverrideFactory<InputT, OutputT, TransformT> { @Override public final Map<PValue, ReplacementOutput> mapOutputs( Map<TupleTag<?>, PValue> outputs, OutputT newOutput) { return ReplacementOutputs.singleton(outputs, newOutput); }  @Override final Map<PValue, ReplacementOutput> mapOutputs(
      Map<TupleTag<?>, PValue> outputs, OutputT newOutput);  }
@Test public void testMapOutputs() { PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3)); PCollection<Integer> output = input.apply("Map", MapElements.via(fn)); PCollection<Integer> reappliedOutput = input.apply("ReMap", MapElements.via(fn)); Map<PValue, ReplacementOutput> replacementMap = factory.mapOutputs(output.expand(), reappliedOutput); assertThat( replacementMap, Matchers.hasEntry( reappliedOutput, ReplacementOutput.of( TaggedPValue.ofExpandedValue(output), TaggedPValue.ofExpandedValue(reappliedOutput)))); }
@Test public void testMapOutputsMultipleOriginalOutputsFails() { PCollection<Integer> input = pipeline.apply(Create.of(1, 2, 3)); PCollection<Integer> output = input.apply("Map", MapElements.via(fn)); PCollection<Integer> reappliedOutput = input.apply("ReMap", MapElements.via(fn)); thrown.expect(IllegalArgumentException.class); factory.mapOutputs( PCollectionList.of(output).and(input).and(reappliedOutput).expand(), reappliedOutput); }