Examples with TestCountingSource org.apache.beam.sdk.runners.dataflow.TestCountingSource used on opensource projects

Example 1 with TestCountingSource

use of org.apache.beam.sdk.runners.dataflow.TestCountingSource in project beam by apache.

the class BoundedReadFromUnboundedSourceTest method testForwardsDisplayData.

@Test
public void testForwardsDisplayData() {
    TestCountingSource src = new TestCountingSource(1234) {

        @Override
        public void populateDisplayData(DisplayData.Builder builder) {
            builder.add(DisplayData.item("foo", "bar"));
        }
    };
    BoundedReadFromUnboundedSource<KV<Integer, Integer>> read = Read.from(src).withMaxNumRecords(5);
    assertThat(DisplayData.from(read), includesDisplayDataFor("source", src));
}

Example 2 with TestCountingSource

use of org.apache.beam.sdk.runners.dataflow.TestCountingSource in project beam by apache.

the class BoundedReadFromUnboundedSourceTest method test.

private void test(boolean dedup, boolean timeBound) throws Exception {
    TestCountingSource source = new TestCountingSource(Integer.MAX_VALUE).withoutSplitting();
    if (dedup) {
        source = source.withDedup();
    }
    PCollection<KV<Integer, Integer>> output = timeBound ? p.apply(Read.from(source).withMaxReadTime(Duration.millis(200))) : p.apply(Read.from(source).withMaxNumRecords(NUM_RECORDS));
    // Because some of the NUM_RECORDS elements read are dupes, the final output
    // will only have output from 0 to n where n < NUM_RECORDS.
    PAssert.that(output).satisfies(new Checker(dedup, timeBound));
    p.run();
}