PTransformReplacements { public static <T> PCollection<T> getSingletonMainInput( AppliedPTransform<? extends PCollection<? extends T>, ?, ?> application) { return getSingletonMainInput( application.getInputs(), application.getTransform().getAdditionalInputs().keySet()); }  static PCollection<T> getSingletonMainInput(
      AppliedPTransform<? extends PCollection<? extends T>, ?, ?> application); static PCollection<T> getSingletonMainOutput(
      AppliedPTransform<?, PCollection<T>, ? extends PTransform<?, PCollection<T>>> transform);  }
@Test public void getMainInputSingleOutputSingleInput() { AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", Collections.singletonMap(new TupleTag<Long>(), mainInput), Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()), pipeline); PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application); assertThat(input, equalTo(mainInput)); }
@Test public void getMainInputSingleOutputSideInputs() { AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", ImmutableMap.<TupleTag<?>, PValue>builder() .put(new TupleTag<Long>(), mainInput) .put(sideInput.getTagInternal(), sideInput.getPCollection()) .build(), Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()).withSideInputs(sideInput), pipeline); PCollection<Long> input = PTransformReplacements.getSingletonMainInput(application); assertThat(input, equalTo(mainInput)); }
@Test public void getMainInputExtraMainInputsThrows() { PCollection<Long> notInParDo = pipeline.apply("otherPCollection", Create.of(1L, 2L, 3L)); ImmutableMap<TupleTag<?>, PValue> inputs = ImmutableMap.<TupleTag<?>, PValue>builder() .putAll(mainInput.expand()) .put(new TupleTag<Long>(), notInParDo) .put(sideInput.getTagInternal(), sideInput.getPCollection()) .build(); AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", inputs, Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()).withSideInputs(sideInput), pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("multiple inputs"); thrown.expectMessage("not additional inputs"); thrown.expectMessage(mainInput.toString()); thrown.expectMessage(notInParDo.toString()); PTransformReplacements.getSingletonMainInput(application); }
@Test public void getMainInputNoMainInputsThrows() { ImmutableMap<TupleTag<?>, PValue> inputs = ImmutableMap.<TupleTag<?>, PValue>builder() .put(sideInput.getTagInternal(), sideInput.getPCollection()) .build(); AppliedPTransform<PCollection<Long>, ?, ?> application = AppliedPTransform.of( "application", inputs, Collections.singletonMap(new TupleTag<Long>(), output), ParDo.of(new TestDoFn()).withSideInputs(sideInput), pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("No main input"); PTransformReplacements.getSingletonMainInput(application); }