/** * Adds the specified elements to the source with timestamp equal to the current watermark. * * @return A {@link TestStream.Builder} like this one that will add the provided elements after * all earlier events have completed. */ @SafeVarargs public final Builder<T> addElements(T element, T... elements) { TimestampedValue<T> firstElement = TimestampedValue.of(element, currentWatermark); @SuppressWarnings({"unchecked", "rawtypes"}) TimestampedValue<T>[] remainingElements = new TimestampedValue[elements.length]; for (int i = 0; i < elements.length; i++) { remainingElements[i] = TimestampedValue.of(elements[i], currentWatermark); } return addElements(firstElement, remainingElements); }
@Override public PCollection<Row> buildIOReader(PBegin begin) { TestStream.Builder<Row> values = TestStream.create( schema, SerializableFunctions.identity(), SerializableFunctions.identity()); for (Pair<Duration, List<Row>> pair : timestampedRows) { values = values.advanceWatermarkTo(new Instant(0).plus(pair.getKey())); for (int i = 0; i < pair.getValue().size(); i++) { values = values.addElements( TimestampedValue.of( pair.getValue().get(i), new Instant(pair.getValue().get(i).getDateTime(timestampField)))); } } return begin .apply( "MockedUnboundedTable_" + COUNTER.incrementAndGet(), values.advanceWatermarkToInfinity()) .setRowSchema(getSchema()); } }
/** * Builds an unbounded {@link PCollection} in {@link Pipeline} set by {@link * #inPipeline(Pipeline)}. * * <p>If timestamp field was set with {@link #withTimestampField(String)} then watermark will be * advanced to the values from that field. */ public PCollection<Row> buildUnbounded() { checkArgument(pipeline != null); checkArgument(rows.size() > 0); if (type == null) { type = rows.get(0).getSchema(); } TestStream.Builder<Row> values = TestStream.create( type, SerializableFunctions.identity(), SerializableFunctions.identity()); for (Row row : rows) { if (timestampField != null) { values = values.advanceWatermarkTo(new Instant(row.getDateTime(timestampField))); } values = values.addElements(row); } return PBegin.in(pipeline).apply("unboundedPCollection", values.advanceWatermarkToInfinity()); } }
/** * Tests that {@link UpdateTeamScoreFn} {@link org.apache.beam.sdk.transforms.DoFn} outputs * correctly per window and per key. */ @Test public void testScoreUpdatesPerWindow() { TestStream<KV<String, GameActionInfo>> createEvents = TestStream.create(KvCoder.of(StringUtf8Coder.of(), AvroCoder.of(GameActionInfo.class))) .advanceWatermarkTo(baseTime) .addElements( event(TestUser.RED_ONE, 50, Duration.standardMinutes(1)), event(TestUser.RED_TWO, 50, Duration.standardMinutes(2)), event(TestUser.RED_ONE, 50, Duration.standardMinutes(3)), event(TestUser.RED_ONE, 60, Duration.standardMinutes(6)), event(TestUser.RED_TWO, 60, Duration.standardMinutes(7))) .advanceWatermarkToInfinity(); Duration teamWindowDuration = Duration.standardMinutes(5); PCollection<KV<String, Integer>> teamScores = p.apply(createEvents) .apply(Window.<KV<String, GameActionInfo>>into(FixedWindows.of(teamWindowDuration))) .apply(ParDo.of(new UpdateTeamScoreFn(100))); String redTeam = TestUser.RED_ONE.getTeam(); String blueTeam = TestUser.BLUE_ONE.getTeam(); IntervalWindow window1 = new IntervalWindow(baseTime, teamWindowDuration); IntervalWindow window2 = new IntervalWindow(window1.end(), teamWindowDuration); PAssert.that(teamScores).inWindow(window1).containsInAnyOrder(KV.of(redTeam, 100)); PAssert.that(teamScores).inWindow(window2).containsInAnyOrder(KV.of(redTeam, 120)); p.run().waitUntilFinish(); }
@Test @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesTestStream.class}) public void testSimpleProcessingTimerTimer() throws Exception { final String timerId = "foo"; DoFn<KV<String, Integer>, Integer> fn = new DoFn<KV<String, Integer>, Integer>() { @TimerId(timerId) private final TimerSpec spec = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); @ProcessElement public void processElement(@TimerId(timerId) Timer timer, OutputReceiver<Integer> r) { timer.offset(Duration.standardSeconds(1)).setRelative(); r.output(3); } @OnTimer(timerId) public void onTimer(TimeDomain timeDomain, OutputReceiver<Integer> r) { if (timeDomain.equals(TimeDomain.PROCESSING_TIME)) { r.output(42); } } }; TestStream<KV<String, Integer>> stream = TestStream.create(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) .addElements(KV.of("hello", 37)) .advanceProcessingTime(Duration.standardSeconds(2)) .advanceWatermarkToInfinity(); PCollection<Integer> output = pipeline.apply(stream).apply(ParDo.of(fn)); PAssert.that(output).containsInAnyOrder(3, 42); pipeline.run(); }
.addElements( event(TestUser.BLUE_ONE, 3, Duration.standardSeconds(3)), event(TestUser.BLUE_ONE, 2, Duration.standardMinutes(1)), .addElements( event(TestUser.RED_ONE, 1, Duration.standardMinutes(4)), event(TestUser.BLUE_ONE, 2, Duration.standardSeconds(270)))
@Test @Category({NeedsRunner.class, UsesTestStream.class}) public void testMultipleStreams() { TestStream<String> stream = TestStream.create(StringUtf8Coder.of()) .addElements("foo", "bar") .advanceWatermarkToInfinity(); TestStream<Integer> other = TestStream.create(VarIntCoder.of()).addElements(1, 2, 3, 4).advanceWatermarkToInfinity(); PCollection<String> createStrings = p.apply("CreateStrings", stream) .apply( "WindowStrings", Window.<String>configure() .triggering(AfterPane.elementCountAtLeast(2)) .withAllowedLateness(Duration.ZERO) .accumulatingFiredPanes()); PAssert.that(createStrings).containsInAnyOrder("foo", "bar"); PCollection<Integer> createInts = p.apply("CreateInts", other) .apply( "WindowInts", Window.<Integer>configure() .triggering(AfterPane.elementCountAtLeast(4)) .withAllowedLateness(Duration.ZERO) .accumulatingFiredPanes()); PAssert.that(createInts).containsInAnyOrder(1, 2, 3, 4); p.run(); }
@Test @Category({NeedsRunner.class, UsesTestStream.class}) public void testProcessingTimeTrigger() { TestStream<Long> source = TestStream.create(VarLongCoder.of()) .addElements( TimestampedValue.of(1L, new Instant(1000L)), TimestampedValue.of(2L, new Instant(2000L))) .advanceProcessingTime(Duration.standardMinutes(12)) .addElements(TimestampedValue.of(3L, new Instant(3000L))) .advanceProcessingTime(Duration.standardMinutes(6)) .advanceWatermarkToInfinity(); PCollection<Long> sum = p.apply(source) .apply( Window.<Long>configure() .triggering( AfterWatermark.pastEndOfWindow() .withEarlyFirings( AfterProcessingTime.pastFirstElementInPane() .plusDelayOf(Duration.standardMinutes(5)))) .accumulatingFiredPanes() .withAllowedLateness(Duration.ZERO)) .apply(Sum.longsGlobally()); PAssert.that(sum).inEarlyGlobalWindowPanes().containsInAnyOrder(3L, 6L); p.run(); }
@Test @Category({NeedsRunner.class, UsesTimersInParDo.class, UsesTestStream.class}) public void testEventTimeTimerUnbounded() throws Exception { final String timerId = "foo"; DoFn<KV<String, Integer>, Integer> fn = new DoFn<KV<String, Integer>, Integer>() { @TimerId(timerId) private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME); @ProcessElement public void processElement(@TimerId(timerId) Timer timer, OutputReceiver<Integer> r) { timer.offset(Duration.standardSeconds(1)).setRelative(); r.output(3); } @OnTimer(timerId) public void onTimer(OutputReceiver<Integer> r) { r.output(42); } }; TestStream<KV<String, Integer>> stream = TestStream.create(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())) .advanceWatermarkTo(new Instant(0)) .addElements(KV.of("hello", 37)) .advanceWatermarkTo(new Instant(0).plus(Duration.standardSeconds(1))) .advanceWatermarkToInfinity(); PCollection<Integer> output = pipeline.apply(stream).apply(ParDo.of(fn)); PAssert.that(output).containsInAnyOrder(3, 42); pipeline.run(); }
/** * Tests that {@link UpdateTeamScoreFn} {@link org.apache.beam.sdk.transforms.DoFn} outputs * correctly for multiple teams. */ @Test public void testScoreUpdatesPerTeam() { TestStream<KV<String, GameActionInfo>> createEvents = TestStream.create(KvCoder.of(StringUtf8Coder.of(), AvroCoder.of(GameActionInfo.class))) .advanceWatermarkTo(baseTime) .addElements( event(TestUser.RED_ONE, 50, Duration.standardSeconds(10)), event(TestUser.RED_TWO, 50, Duration.standardSeconds(20)), event(TestUser.BLUE_ONE, 70, Duration.standardSeconds(30)), event(TestUser.BLUE_TWO, 80, Duration.standardSeconds(40)), event(TestUser.BLUE_TWO, 50, Duration.standardSeconds(50))) .advanceWatermarkToInfinity(); PCollection<KV<String, Integer>> teamScores = p.apply(createEvents).apply(ParDo.of(new UpdateTeamScoreFn(100))); String redTeam = TestUser.RED_ONE.getTeam(); String blueTeam = TestUser.BLUE_ONE.getTeam(); PAssert.that(teamScores) .inWindow(GlobalWindow.INSTANCE) .containsInAnyOrder(KV.of(redTeam, 100), KV.of(blueTeam, 150), KV.of(blueTeam, 200)); p.run().waitUntilFinish(); }
@Test @Category({NeedsRunner.class, UsesTestStream.class}) public void testElementsAtAlmostPositiveInfinity() { Instant endOfGlobalWindow = GlobalWindow.INSTANCE.maxTimestamp(); TestStream<String> stream = TestStream.create(StringUtf8Coder.of()) .addElements( TimestampedValue.of("foo", endOfGlobalWindow), TimestampedValue.of("bar", endOfGlobalWindow)) .advanceWatermarkToInfinity(); FixedWindows windows = FixedWindows.of(Duration.standardHours(6)); PCollection<String> windowedValues = p.apply(stream) .apply(Window.into(windows)) .apply(WithKeys.of(1)) .apply(GroupByKey.create()) .apply(Values.create()) .apply(Flatten.iterables()); PAssert.that(windowedValues) .inWindow(windows.assignWindow(endOfGlobalWindow)) .containsInAnyOrder("foo", "bar"); p.run(); }
@Test @Category({NeedsRunner.class, UsesTestStream.class}) public void streamingWrites() throws Exception { TestStream<Mutation> testStream = TestStream.create(SerializableCoder.of(Mutation.class)) .addElements(m(1L), m(2L)) .advanceProcessingTime(Duration.standardMinutes(1)) .addElements(m(3L), m(4L)) .advanceProcessingTime(Duration.standardMinutes(1)) .addElements(m(5L), m(6L)) .advanceWatermarkToInfinity(); pipeline .apply(testStream) .apply( SpannerIO.write() .withProjectId("test-project") .withInstanceId("test-instance") .withDatabaseId("test-database") .withServiceFactory(serviceFactory)); pipeline.run(); verifyBatches(batch(m(1L), m(2L)), batch(m(3L), m(4L)), batch(m(5L), m(6L))); }
/** * Tests that {@link UpdateTeamScoreFn} {@link org.apache.beam.sdk.transforms.DoFn} outputs * correctly for one team. */ @Test public void testScoreUpdatesOneTeam() { TestStream<KV<String, GameActionInfo>> createEvents = TestStream.create(KvCoder.of(StringUtf8Coder.of(), AvroCoder.of(GameActionInfo.class))) .advanceWatermarkTo(baseTime) .addElements( event(TestUser.RED_TWO, 99, Duration.standardSeconds(10)), event(TestUser.RED_ONE, 1, Duration.standardSeconds(20)), event(TestUser.RED_ONE, 0, Duration.standardSeconds(30)), event(TestUser.RED_TWO, 100, Duration.standardSeconds(40)), event(TestUser.RED_TWO, 201, Duration.standardSeconds(50))) .advanceWatermarkToInfinity(); PCollection<KV<String, Integer>> teamScores = p.apply(createEvents).apply(ParDo.of(new UpdateTeamScoreFn(100))); String redTeam = TestUser.RED_ONE.getTeam(); PAssert.that(teamScores) .inWindow(GlobalWindow.INSTANCE) .containsInAnyOrder(KV.of(redTeam, 100), KV.of(redTeam, 200), KV.of(redTeam, 401)); p.run().waitUntilFinish(); }
@Test @Category({ValidatesRunner.class, UsesTestStream.class}) public void testReshuffleWithTimestampsStreaming() { TestStream<Long> stream = TestStream.create(VarLongCoder.of()) .advanceWatermarkTo(new Instant(0L).plus(Duration.standardDays(48L))) .addElements( TimestampedValue.of(0L, new Instant(0L)), TimestampedValue.of(1L, new Instant(0L).plus(Duration.standardDays(48L))), TimestampedValue.of( 2L, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(Duration.standardDays(48L)))) .advanceWatermarkToInfinity(); PCollection<KV<String, Long>> input = pipeline .apply(stream) .apply(WithKeys.of("")) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(10L)))); PCollection<KV<String, Long>> reshuffled = input.apply(Reshuffle.of()); PAssert.that(reshuffled.apply(Values.create())).containsInAnyOrder(0L, 1L, 2L); pipeline.run(); } }
@Test @Category(NeedsRunner.class) public void globalWindowNoKeys() { PCollection<ValueInSingleWindow<String>> result = pipeline .apply( TestStream.create(StringUtf8Coder.of()) .addElements(TimestampedValue.of("dei", new Instant(123L))) .advanceWatermarkToInfinity()) .apply(Reify.windows()); PAssert.that(result) .containsInAnyOrder( ValueInSingleWindow.of( "dei", new Instant(123L), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)); pipeline.run(); }
private PCollection<Row> prepareUnboundedPCollection2() { TestStream.Builder<Row> values = TestStream.create( schemaInTableA, SerializableFunctions.identity(), SerializableFunctions.identity()); Row row = rowsInTableA.get(0); values = values.advanceWatermarkTo(new Instant(row.getDateTime("f_timestamp"))); values = values.addElements(row); return PBegin.in(pipeline) .apply("unboundedInput2", values.advanceWatermarkToInfinity()) .apply( "unboundedInput2.fixedWindow1year", Window.into(FixedWindows.of(Duration.standardDays(365)))); } }
private PCollection<Row> prepareUnboundedPCollection1() { TestStream.Builder<Row> values = TestStream.create( schemaInTableA, SerializableFunctions.identity(), SerializableFunctions.identity()); for (Row row : rowsInTableA) { values = values.advanceWatermarkTo(new Instant(row.getDateTime("f_timestamp"))); values = values.addElements(row); } return PBegin.in(pipeline) .apply("unboundedInput1", values.advanceWatermarkToInfinity()) .apply( "unboundedInput1.fixedWindow1year", Window.into(FixedWindows.of(Duration.standardDays(365)))); }
@Parameters(name = "{index}: {0}") public static Iterable<TestStream<?>> data() { return ImmutableList.of( TestStream.create(VarIntCoder.of()).advanceWatermarkToInfinity(), TestStream.create(VarIntCoder.of()) .advanceWatermarkTo(new Instant(42)) .advanceWatermarkToInfinity(), TestStream.create(VarIntCoder.of()) .addElements(TimestampedValue.of(3, new Instant(17))) .advanceWatermarkToInfinity(), TestStream.create(StringUtf8Coder.of()) .advanceProcessingTime(Duration.millis(82)) .advanceWatermarkToInfinity()); }
@Test @Category({ValidatesRunner.class, UsesTestStream.class}) public void duplicateTimerSetting() { TestStream<KV<String, String>> stream = TestStream.create(KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of())) .addElements(KV.of("key1", "v1")) .advanceWatermarkToInfinity(); PCollection<String> result = pipeline.apply(stream).apply(ParDo.of(new TwoTimerDoFn())); PAssert.that(result).containsInAnyOrder("It works"); pipeline.run().waitUntilFinish(); } }
@Test public void testElementAtPositiveInfinityThrows() { Builder<Integer> stream = TestStream.create(VarIntCoder.of()) .addElements(TimestampedValue.of(-1, BoundedWindow.TIMESTAMP_MAX_VALUE.minus(1L))); thrown.expect(IllegalArgumentException.class); stream.addElements(TimestampedValue.of(1, BoundedWindow.TIMESTAMP_MAX_VALUE)); }