use of org.apache.beam.sdk.state.TimerSpec in project beam by apache.
the class DoFnInvokersTest method testOnTimerWithWindow.
@Test
public void testOnTimerWithWindow() throws Exception {
final String timerId = "my-timer-id";
final IntervalWindow testWindow = new IntervalWindow(new Instant(0), new Instant(15));
when(mockArgumentProvider.window()).thenReturn(testWindow);
class SimpleTimerDoFn extends DoFn<String, String> {
public IntervalWindow window = null;
@TimerId(timerId)
private final TimerSpec myTimer = TimerSpecs.timer(TimeDomain.PROCESSING_TIME);
@ProcessElement
public void process(ProcessContext c) {
}
@OnTimer(timerId)
public void onMyTimer(IntervalWindow w) {
window = w;
}
}
SimpleTimerDoFn fn = new SimpleTimerDoFn();
DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn);
invoker.invokeOnTimer(timerId, mockArgumentProvider);
assertThat(fn.window, equalTo(testWindow));
}
use of org.apache.beam.sdk.state.TimerSpec in project beam by apache.
the class DoFnInvokersTest method testDoFnWithTimer.
/**
* Tests that the generated {@link DoFnInvoker} passes the timer parameter that it
* should.
*/
@Test
public void testDoFnWithTimer() throws Exception {
Timer mockTimer = mock(Timer.class);
final String timerId = "my-timer-id-here";
when(mockArgumentProvider.timer(timerId)).thenReturn(mockTimer);
class MockFn extends DoFn<String, String> {
@TimerId(timerId)
private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);
@ProcessElement
public void processElement(ProcessContext c, @TimerId(timerId) Timer timer) throws Exception {
}
@OnTimer(timerId)
public void onTimer() {
}
}
MockFn fn = mock(MockFn.class);
invokeProcessElement(fn);
verify(fn).processElement(mockProcessContext, mockTimer);
}
use of org.apache.beam.sdk.state.TimerSpec in project beam by apache.
the class ParDoTest method testEventTimeTimerUnbounded.
@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(ProcessContext context, @TimerId(timerId) Timer timer) {
timer.offset(Duration.standardSeconds(1)).setRelative();
context.output(3);
}
@OnTimer(timerId)
public void onTimer(OnTimerContext context) {
context.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();
}
use of org.apache.beam.sdk.state.TimerSpec in project beam by apache.
the class ParDoTest method testEventTimeTimerBounded.
/**
* Tests that an event time timer fires and results in supplementary output.
*
* <p>This test relies on two properties:
*
* <ol>
* <li>A timer that is set on time should always get a chance to fire. For this to be true, timers
* per-key-and-window must be delivered in order so the timer is not wiped out until the
* window is expired by the runner.
* <li>A {@link Create} transform sends its elements on time, and later advances the watermark to
* infinity
* </ol>
*
* <p>Note that {@link TestStream} is not applicable because it requires very special runner hooks
* and is only supported by the direct runner.
*/
@Test
@Category({ ValidatesRunner.class, UsesTimersInParDo.class })
public void testEventTimeTimerBounded() 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(ProcessContext context, @TimerId(timerId) Timer timer) {
timer.offset(Duration.standardSeconds(1)).setRelative();
context.output(3);
}
@OnTimer(timerId)
public void onTimer(OnTimerContext context) {
context.output(42);
}
};
PCollection<Integer> output = pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
PAssert.that(output).containsInAnyOrder(3, 42);
pipeline.run();
}
use of org.apache.beam.sdk.state.TimerSpec in project beam by apache.
the class ParDoTest method testAbsoluteProcessingTimeTimerRejected.
@Test
@Category({ ValidatesRunner.class, UsesTimersInParDo.class })
public void testAbsoluteProcessingTimeTimerRejected() 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(ProcessContext context, @TimerId(timerId) Timer timer) {
timer.set(new Instant(0));
}
@OnTimer(timerId)
public void onTimer(OnTimerContext context) {
}
};
PCollection<Integer> output = pipeline.apply(Create.of(KV.of("hello", 37))).apply(ParDo.of(fn));
thrown.expect(RuntimeException.class);
// Note that runners can reasonably vary their message - this matcher should be flexible
// and can be evolved.
thrown.expectMessage("relative timers");
thrown.expectMessage("processing time");
pipeline.run();
}
Aggregations