TransformInputs { public static Collection<PValue> nonAdditionalInputs(AppliedPTransform<?, ?, ?> application) { ImmutableList.Builder<PValue> mainInputs = ImmutableList.builder(); PTransform<?, ?> transform = application.getTransform(); for (Map.Entry<TupleTag<?>, PValue> input : application.getInputs().entrySet()) { if (!transform.getAdditionalInputs().containsKey(input.getKey())) { mainInputs.add(input.getValue()); } } checkArgument( !mainInputs.build().isEmpty() || application.getInputs().isEmpty(), "Expected at least one main input if any inputs exist"); return mainInputs.build(); }  static Collection<PValue> nonAdditionalInputs(AppliedPTransform<?, ?, ?> application);  }
@Test public void nonAdditionalInputsWithNoInputSucceeds() { AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "input-free", Collections.emptyMap(), Collections.emptyMap(), new TestTransform(), pipeline); assertThat(TransformInputs.nonAdditionalInputs(transform), Matchers.empty()); }
@Test public void nonAdditionalInputsWithOneMainInputSucceeds() { PCollection<Long> input = pipeline.apply(GenerateSequence.from(1L)); AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "input-single", Collections.singletonMap(new TupleTag<Long>() {}, input), Collections.emptyMap(), new TestTransform(), pipeline); assertThat(TransformInputs.nonAdditionalInputs(transform), Matchers.containsInAnyOrder(input)); }
@Test public void nonAdditionalInputsWithMultipleNonAdditionalInputsSucceeds() { Map<TupleTag<?>, PValue> allInputs = new HashMap<>(); PCollection<Integer> mainInts = pipeline.apply("MainInput", Create.of(12, 3)); allInputs.put(new TupleTag<Integer>() {}, mainInts); PCollection<Void> voids = pipeline.apply("VoidInput", Create.empty(VoidCoder.of())); allInputs.put(new TupleTag<Void>() {}, voids); AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "additional-free", allInputs, Collections.emptyMap(), new TestTransform(), pipeline); assertThat( TransformInputs.nonAdditionalInputs(transform), Matchers.containsInAnyOrder(voids, mainInts)); }
@Test public void nonAdditionalInputsWithAdditionalInputsSucceeds() { Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>(); additionalInputs.put(new TupleTag<String>() {}, pipeline.apply(Create.of("1, 2", "3"))); additionalInputs.put(new TupleTag<Long>() {}, pipeline.apply(GenerateSequence.from(3L))); Map<TupleTag<?>, PValue> allInputs = new HashMap<>(); PCollection<Integer> mainInts = pipeline.apply("MainInput", Create.of(12, 3)); allInputs.put(new TupleTag<Integer>() {}, mainInts); PCollection<Void> voids = pipeline.apply("VoidInput", Create.empty(VoidCoder.of())); allInputs.put(new TupleTag<Void>() {}, voids); allInputs.putAll(additionalInputs); AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "additional", allInputs, Collections.emptyMap(), new TestTransform(additionalInputs), pipeline); assertThat( TransformInputs.nonAdditionalInputs(transform), Matchers.containsInAnyOrder(mainInts, voids)); }
@Test public void nonAdditionalInputsWithOnlyAdditionalInputsThrows() { Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>(); additionalInputs.put(new TupleTag<String>() {}, pipeline.apply(Create.of("1, 2", "3"))); additionalInputs.put(new TupleTag<Long>() {}, pipeline.apply(GenerateSequence.from(3L))); AppliedPTransform<PInput, POutput, TestTransform> transform = AppliedPTransform.of( "additional-only", additionalInputs, Collections.emptyMap(), new TestTransform(additionalInputs), pipeline); thrown.expect(IllegalArgumentException.class); thrown.expectMessage("at least one"); TransformInputs.nonAdditionalInputs(transform); }