UnconsumedReads { public static void ensureAllReadsConsumed(Pipeline pipeline) { final Set<PCollection<?>> unconsumed = new HashSet<>(); pipeline.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { unconsumed.removeAll(node.getInputs().values()); } @Override public void visitValue(PValue value, Node producer) { String urn = PTransformTranslation.urnForTransformOrNull(producer.getTransform()); if (PTransformTranslation.READ_TRANSFORM_URN.equals(urn)) { unconsumed.add((PCollection<?>) value); } } }); int i = 0; for (PCollection<?> unconsumedPCollection : unconsumed) { consume(unconsumedPCollection, i); i++; } }  static void ensureAllReadsConsumed(Pipeline pipeline);  }
@Test public void matcherProducesUnconsumedValueBoundedRead() { Bounded<Long> transform = Read.from(CountingSource.upTo(20L)); pipeline.apply(transform); UnconsumedReads.ensureAllReadsConsumed(pipeline); validateConsumed(); }
@Test public void matcherProducesUnconsumedValueUnboundedRead() { Unbounded<Long> transform = Read.from(CountingSource.unbounded()); pipeline.apply(transform); UnconsumedReads.ensureAllReadsConsumed(pipeline); validateConsumed(); }
@Test public void doesNotConsumeAlreadyConsumedRead() { Unbounded<Long> transform = Read.from(CountingSource.unbounded()); final PCollection<Long> output = pipeline.apply(transform); final Flatten.PCollections<Long> consumer = Flatten.pCollections(); PCollectionList.of(output).apply(consumer); UnconsumedReads.ensureAllReadsConsumed(pipeline); pipeline.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { if (node.getInputs().values().contains(output)) { assertThat(node.getTransform(), Matchers.is(consumer)); } } }); }