DoFnInvokers { public static <InputT, OutputT> DoFnInvoker<InputT, OutputT> invokerFor( DoFn<InputT, OutputT> fn) { return ByteBuddyDoFnInvokerFactory.only().newByteBuddyInvoker(fn); } private  DoFnInvokers(); static DoFnInvoker<InputT, OutputT> invokerFor(
      DoFn<InputT, OutputT> fn); static DoFnInvoker<InputT, OutputT> tryInvokeSetupFor(
      DoFn<InputT, OutputT> fn);  }
@Test public void testDoFnInvokersReused() throws Exception { IdentityParent fn1 = new IdentityParent(); IdentityParent fn2 = new IdentityParent(); assertSame( "Invoker classes should only be generated once for each type", DoFnInvokers.invokerFor(fn1).getClass(), DoFnInvokers.invokerFor(fn2).getClass()); }
@Test public void testOnWindowExpirationWithNoParam() throws Exception { class MockFn extends DoFn<String, String> { @ProcessElement public void process(ProcessContext c) {} @OnWindowExpiration public void onWindowExpiration() {} } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeOnWindowExpiration(mockArgumentProvider); verify(fn).onWindowExpiration(); }
@Test public void testOnWindowExpirationWithParam() { class MockFn extends DoFn<String, String> { @ProcessElement public void process(ProcessContext c) {} @OnWindowExpiration public void onWindowExpiration(BoundedWindow window) {} } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeOnWindowExpiration(mockArgumentProvider); verify(fn).onWindowExpiration(mockWindow); }
@Test public void testDoFnWithStartBundleSetupTeardown() throws Exception { when(mockArgumentProvider.startBundleContext(any(DoFn.class))) .thenReturn(mockStartBundleContext); when(mockArgumentProvider.finishBundleContext(any(DoFn.class))) .thenReturn(mockFinishBundleContext); class MockFn extends DoFn<String, String> { @ProcessElement public void processElement(ProcessContext c) {} @StartBundle public void startBundle(StartBundleContext c) {} @FinishBundle public void finishBundle(FinishBundleContext c) {} @Setup public void before() {} @Teardown public void after() {} } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeSetup(); invoker.invokeStartBundle(mockArgumentProvider); invoker.invokeFinishBundle(mockArgumentProvider); invoker.invokeTeardown(); verify(fn).before(); verify(fn).startBundle(mockStartBundleContext); verify(fn).finishBundle(mockFinishBundleContext); verify(fn).after(); }
@Test public void testSplittableDoFnWithAllMethods() throws Exception { MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); final SomeRestrictionTracker tracker = mock(SomeRestrictionTracker.class); final SomeRestrictionCoder coder = mock(SomeRestrictionCoder.class); final InstantCoder watermarkEstimatorStateCoder = InstantCoder.of(); final Instant watermarkEstimatorState = Instant.now(); final WatermarkEstimator<Instant> watermarkEstimator = new WatermarkEstimators.Manual(watermarkEstimatorState); SomeRestriction restriction = new SomeRestriction(); final SomeRestriction part1 = new SomeRestriction(); final SomeRestriction part2 = new SomeRestriction(); final SomeRestriction part3 = new SomeRestriction(); when(fn.getRestrictionCoder()).thenReturn(coder); when(fn.getWatermarkEstimatorStateCoder()).thenReturn(watermarkEstimatorStateCoder); when(fn.getInitialRestriction(mockElement)).thenReturn(restriction); doAnswer( AdditionalAnswers.delegatesTo( new MockFn() { @DoFn.SplitRestriction @Override public void splitRestriction( @Element String element, @Restriction SomeRestriction restriction, DoFn.OutputReceiver<SomeRestriction> receiver) { receiver.output(part1); receiver.output(part2); receiver.output(part3); } })) .when(fn) .splitRestriction(eq(mockElement), same(restriction), any()); when(fn.getInitialWatermarkEstimatorState()).thenReturn(watermarkEstimatorState); when(fn.newTracker(restriction)).thenReturn(tracker); when(fn.newWatermarkEstimator(watermarkEstimatorState)).thenReturn(watermarkEstimator); when(fn.processElement(mockProcessContext, tracker, watermarkEstimator)).thenReturn(resume()); assertEquals(coder, invoker.invokeGetRestrictionCoder(CoderRegistry.createDefault())); assertEquals( watermarkEstimatorStateCoder, invoker.invokeGetWatermarkEstimatorStateCoder(CoderRegistry.createDefault())); assertEquals( restriction, invoker.invokeGetInitialRestriction( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return mockElement; } })); List<SomeRestriction> outputs = new ArrayList<>(); invoker.invokeSplitRestriction( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return mockElement; } @Override public Object restriction() { return restriction; } @Override public OutputReceiver outputReceiver(DoFn doFn) { return new OutputReceiver<SomeRestriction>() { @Override public void output(SomeRestriction output) { outputs.add(output); } @Override public void outputWithTimestamp(SomeRestriction output, Instant timestamp) { fail("Unexpected output with timestamp"); } }; } }); assertEquals(Arrays.asList(part1, part2, part3), outputs); assertEquals( watermarkEstimatorState, invoker.invokeGetInitialWatermarkEstimatorState(new FakeArgumentProvider<>())); assertEquals( tracker, invoker.invokeNewTracker( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return mockElement; } @Override public Object restriction() { return restriction; } })); assertEquals( watermarkEstimator, invoker.invokeNewWatermarkEstimator( new FakeArgumentProvider<String, String>() { @Override public Object watermarkEstimatorState() { return watermarkEstimatorState; } })); assertEquals( resume(), invoker.invokeProcessElement( new FakeArgumentProvider<String, String>() { @Override public DoFn<String, String>.ProcessContext processContext(DoFn<String, String> fn) { return mockProcessContext; } @Override public RestrictionTracker<?, ?> restrictionTracker() { return tracker; } @Override public WatermarkEstimator<?> watermarkEstimator() { return watermarkEstimator; } })); }
@Test public void testSplittableDoFnWithHasDefaultMethods() throws Exception { class MockFn extends DoFn<String, String> { @ProcessElement public void processElement( ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {} @GetInitialRestriction public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) { return null; } @GetInitialWatermarkEstimatorState public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() { return null; } } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); CoderRegistry coderRegistry = CoderRegistry.createDefault(); coderRegistry.registerCoderProvider( CoderProviders.fromStaticMethods( RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class)); coderRegistry.registerCoderForClass( WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator()); assertThat( invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class)); assertThat( invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class)); invoker.invokeSplitRestriction( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return "blah"; } @Override public Object restriction() { return "foo"; } @Override public OutputReceiver<String> outputReceiver(DoFn<String, String> doFn) { return new DoFn.OutputReceiver<String>() { private boolean invoked; @Override public void output(String output) { assertFalse(invoked); invoked = true; assertEquals("foo", output); } @Override public void outputWithTimestamp(String output, Instant instant) { assertFalse(invoked); invoked = true; assertEquals("foo", output); } }; } }); assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider)); assertThat( invoker.invokeNewTracker( new FakeArgumentProvider<String, String>() { @Override public Object restriction() { return new RestrictionWithBoundedDefaultTracker(); } }), instanceOf(BoundedDefaultTracker.class)); assertThat( invoker.invokeNewWatermarkEstimator( new FakeArgumentProvider<String, String>() { @Override public Object watermarkEstimatorState() { return new WatermarkEstimatorStateWithDefaultWatermarkEstimator(); } }), instanceOf(DefaultWatermarkEstimator.class)); }
@Test public void testTruncateFnWithHasDefaultMethodsWhenBounded() throws Exception { class BoundedMockFn extends DoFn<String, String> { @ProcessElement public void processElement( ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {} @GetInitialRestriction public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) { return null; } @GetInitialWatermarkEstimatorState public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() { return null; } } BoundedMockFn fn = mock(BoundedMockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); CoderRegistry coderRegistry = CoderRegistry.createDefault(); coderRegistry.registerCoderProvider( CoderProviders.fromStaticMethods( RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class)); coderRegistry.registerCoderForClass( WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator()); assertThat( invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class)); assertThat( invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class)); RestrictionTracker tracker = invoker.invokeNewTracker( new FakeArgumentProvider<String, String>() { @Override public Object restriction() { return new RestrictionWithBoundedDefaultTracker(); } }); assertThat(tracker, instanceOf(BoundedDefaultTracker.class)); TruncateResult<?> result = invoker.invokeTruncateRestriction( new FakeArgumentProvider<String, String>() { @Override public RestrictionTracker restrictionTracker() { return tracker; } @Override public String element(DoFn<String, String> doFn) { return "blah"; } @Override public Object restriction() { return "foo"; } }); assertEquals("foo", result.getTruncatedRestriction()); assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider)); assertThat( invoker.invokeNewWatermarkEstimator( new FakeArgumentProvider<String, String>() { @Override public Object watermarkEstimatorState() { return new WatermarkEstimatorStateWithDefaultWatermarkEstimator(); } }), instanceOf(DefaultWatermarkEstimator.class)); }
@Test public void testTruncateFnWithHasDefaultMethodsWhenUnbounded() throws Exception { class UnboundedMockFn extends DoFn<String, String> { @ProcessElement public void processElement( ProcessContext c, RestrictionTracker<RestrictionWithUnboundedDefaultTracker, Void> tracker, WatermarkEstimator<WatermarkEstimatorStateWithDefaultWatermarkEstimator> watermarkEstimator) {} @GetInitialRestriction public RestrictionWithUnboundedDefaultTracker getInitialRestriction(@Element String element) { return null; } @GetInitialWatermarkEstimatorState public WatermarkEstimatorStateWithDefaultWatermarkEstimator getInitialWatermarkEstimatorState() { return null; } } UnboundedMockFn fn = mock(UnboundedMockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); CoderRegistry coderRegistry = CoderRegistry.createDefault(); coderRegistry.registerCoderProvider( CoderProviders.fromStaticMethods( RestrictionWithUnboundedDefaultTracker.class, CoderForDefaultTracker.class)); coderRegistry.registerCoderForClass( WatermarkEstimatorStateWithDefaultWatermarkEstimator.class, new CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator()); assertThat( invoker.<RestrictionWithBoundedDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class)); assertThat( invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry), instanceOf(CoderForWatermarkEstimatorStateWithDefaultWatermarkEstimator.class)); RestrictionTracker tracker = invoker.invokeNewTracker( new FakeArgumentProvider<String, String>() { @Override public Object restriction() { return new RestrictionWithUnboundedDefaultTracker(); } }); assertThat(tracker, instanceOf(UnboundedDefaultTracker.class)); TruncateResult<?> result = invoker.invokeTruncateRestriction( new FakeArgumentProvider<String, String>() { @Override public RestrictionTracker restrictionTracker() { return tracker; } @Override public String element(DoFn<String, String> doFn) { return "blah"; } @Override public Object restriction() { return "foo"; } }); assertNull(result); assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider)); assertThat( invoker.invokeNewWatermarkEstimator( new FakeArgumentProvider<String, String>() { @Override public Object watermarkEstimatorState() { return new WatermarkEstimatorStateWithDefaultWatermarkEstimator(); } }), instanceOf(DefaultWatermarkEstimator.class)); }
@Test public void testDefaultWatermarkEstimatorStateAndCoder() throws Exception { class MockFn extends DoFn<String, String> { @ProcessElement public void processElement( ProcessContext c, RestrictionTracker<RestrictionWithBoundedDefaultTracker, Void> tracker) {} @GetInitialRestriction public RestrictionWithBoundedDefaultTracker getInitialRestriction(@Element String element) { return null; } } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); CoderRegistry coderRegistry = CoderRegistry.createDefault(); coderRegistry.registerCoderProvider( CoderProviders.fromStaticMethods( RestrictionWithBoundedDefaultTracker.class, CoderForDefaultTracker.class)); assertEquals(VoidCoder.of(), invoker.invokeGetWatermarkEstimatorStateCoder(coderRegistry)); assertNull(invoker.invokeGetInitialWatermarkEstimatorState(new FakeArgumentProvider<>())); }
@Test public void testProcessElementException() throws Exception { DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) { throw new IllegalArgumentException("bogus"); } }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeProcessElement( new FakeArgumentProvider<Integer, Integer>() { @Override public DoFn<Integer, Integer>.ProcessContext processContext(DoFn<Integer, Integer> fn) { return null; } }); }
@Test public void testProcessElementExceptionWithReturn() throws Exception { thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @ProcessElement public ProcessContinuation processElement( @SuppressWarnings("unused") ProcessContext c, RestrictionTracker<SomeRestriction, Void> tracker) { throw new IllegalArgumentException("bogus"); } @GetInitialRestriction public SomeRestriction getInitialRestriction(@Element Integer element) { return null; } @NewTracker public SomeRestrictionTracker newTracker(@Restriction SomeRestriction restriction) { return null; } }) .invokeProcessElement( new FakeArgumentProvider<Integer, Integer>() { @Override public DoFn.ProcessContext processContext(DoFn<Integer, Integer> doFn) { return null; } @Override public RestrictionTracker<?, ?> restrictionTracker() { return null; } }); }
@Test public void testStartBundleException() throws Exception { DoFnInvoker.ArgumentProvider<Integer, Integer> mockArguments = mock(DoFnInvoker.ArgumentProvider.class); when(mockArguments.startBundleContext(any(DoFn.class))).thenReturn(null); DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @StartBundle public void startBundle(@SuppressWarnings("unused") StartBundleContext c) { throw new IllegalArgumentException("bogus"); } @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) {} }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeStartBundle(mockArguments); }
@Test public void testFinishBundleException() throws Exception { DoFnInvoker.ArgumentProvider<Integer, Integer> mockArguments = mock(DoFnInvoker.ArgumentProvider.class); when(mockArguments.finishBundleContext(any(DoFn.class))).thenReturn(null); DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @FinishBundle public void finishBundle(@SuppressWarnings("unused") FinishBundleContext c) { throw new IllegalArgumentException("bogus"); } @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) {} }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeFinishBundle(mockArguments); }
@Test public void testOnTimerHelloWord() throws Exception { final String timerId = "my-timer-id"; class SimpleTimerDoFn extends DoFn<String, String> { public String status = "not yet"; @TimerId(timerId) private final TimerSpec myTimer = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); @ProcessElement public void process(ProcessContext c) {} @OnTimer(timerId) public void onMyTimer() { status = "OK now"; } } SimpleTimerDoFn fn = new SimpleTimerDoFn(); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeOnTimer(TimerDeclaration.PREFIX + timerId, "", mockArgumentProvider); assertThat(fn.status, equalTo("OK now")); }
@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(TimerDeclaration.PREFIX + timerId, "", mockArgumentProvider); assertThat(fn.window, equalTo(testWindow)); }
@Test public void testStableName() { DoFnInvoker<Void, Void> invoker = DoFnInvokers.invokerFor(new StableNameTestDoFn()); assertThat( invoker.getClass().getName(), equalTo( String.format( "%s$%s", StableNameTestDoFn.class.getName(), DoFnInvoker.class.getSimpleName()))); }
@Test public void testBundleFinalizer() { class BundleFinalizerDoFn extends DoFn<String, String> { @ProcessElement public void processElement(BundleFinalizer bundleFinalizer) { bundleFinalizer.afterBundleCommit(Instant.ofEpochSecond(42L), null); } } BundleFinalizer mockBundleFinalizer = mock(BundleFinalizer.class); when(mockArgumentProvider.bundleFinalizer()).thenReturn(mockBundleFinalizer); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(new BundleFinalizerDoFn()); invoker.invokeProcessElement(mockArgumentProvider); verify(mockBundleFinalizer).afterBundleCommit(eq(Instant.ofEpochSecond(42L)), eq(null)); }
@Test public void testDoFnWithStartBundleSetupTeardown() throws Exception { class MockFn extends DoFn<String, String> { @ProcessElement public void processElement(ProcessContext c) {} @StartBundle public void startBundle(StartBundleContext c) {} @FinishBundle public void finishBundle(FinishBundleContext c) {} @Setup public void before() {} @Teardown public void after() {} } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeSetup(); invoker.invokeStartBundle(mockStartBundleContext); invoker.invokeFinishBundle(mockFinishBundleContext); invoker.invokeTeardown(); verify(fn).before(); verify(fn).startBundle(mockStartBundleContext); verify(fn).finishBundle(mockFinishBundleContext); verify(fn).after(); }
@Test public void testSplittableDoFnWithAllMethods() throws Exception { MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); final SomeRestrictionTracker tracker = mock(SomeRestrictionTracker.class); final SomeRestrictionCoder coder = mock(SomeRestrictionCoder.class); SomeRestriction restriction = new SomeRestriction(); final SomeRestriction part1 = new SomeRestriction(); final SomeRestriction part2 = new SomeRestriction(); final SomeRestriction part3 = new SomeRestriction(); when(fn.getRestrictionCoder()).thenReturn(coder); when(fn.getInitialRestriction(mockElement)).thenReturn(restriction); doAnswer( AdditionalAnswers.delegatesTo( new MockFn() { @DoFn.SplitRestriction @Override public void splitRestriction( @Element String element, @Restriction SomeRestriction restriction, DoFn.OutputReceiver<SomeRestriction> receiver) { receiver.output(part1); receiver.output(part2); receiver.output(part3); } })) .when(fn) .splitRestriction(eq(mockElement), same(restriction), Mockito.any()); when(fn.newTracker(restriction)).thenReturn(tracker); when(fn.processElement(mockProcessContext, tracker)).thenReturn(resume()); assertEquals(coder, invoker.invokeGetRestrictionCoder(CoderRegistry.createDefault())); assertEquals( restriction, invoker.invokeGetInitialRestriction( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return mockElement; } })); List<SomeRestriction> outputs = new ArrayList<>(); invoker.invokeSplitRestriction( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return mockElement; } @Override public Object restriction() { return restriction; } @Override public OutputReceiver outputReceiver(DoFn doFn) { return new OutputReceiver<SomeRestriction>() { @Override public void output(SomeRestriction output) { outputs.add(output); } @Override public void outputWithTimestamp(SomeRestriction output, Instant timestamp) { fail("Unexpected output with timestamp"); } }; } }); assertEquals(Arrays.asList(part1, part2, part3), outputs); assertEquals( tracker, invoker.invokeNewTracker( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return mockElement; } @Override public Object restriction() { return restriction; } })); assertEquals( resume(), invoker.invokeProcessElement( new FakeArgumentProvider<String, String>() { @Override public DoFn<String, String>.ProcessContext processContext(DoFn<String, String> fn) { return mockProcessContext; } @Override public RestrictionTracker<?, ?> restrictionTracker() { return tracker; } })); }
@Test public void testSplittableDoFnDefaultMethods() throws Exception { class MockFn extends DoFn<String, String> { @ProcessElement public void processElement( ProcessContext c, RestrictionTracker<RestrictionWithDefaultTracker, Void> tracker) {} @GetInitialRestriction public RestrictionWithDefaultTracker getInitialRestriction(@Element String element) { return null; } } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); CoderRegistry coderRegistry = CoderRegistry.createDefault(); coderRegistry.registerCoderProvider( CoderProviders.fromStaticMethods( RestrictionWithDefaultTracker.class, CoderForDefaultTracker.class)); assertThat( invoker.<RestrictionWithDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class)); invoker.invokeSplitRestriction( new FakeArgumentProvider<String, String>() { @Override public String element(DoFn<String, String> doFn) { return "blah"; } @Override public Object restriction() { return "foo"; } @Override public OutputReceiver<String> outputReceiver(DoFn<String, String> doFn) { return new DoFn.OutputReceiver<String>() { private boolean invoked; @Override public void output(String output) { assertFalse(invoked); invoked = true; assertEquals("foo", output); } @Override public void outputWithTimestamp(String output, Instant instant) { assertFalse(invoked); invoked = true; assertEquals("foo", output); } }; } }); assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider)); assertThat( invoker.invokeNewTracker( new FakeArgumentProvider<String, String>() { @Override public Object restriction() { return new RestrictionWithDefaultTracker(); } }), instanceOf(DefaultTracker.class)); }
@Test public void testStartBundleException() throws Exception { DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @StartBundle public void startBundle(@SuppressWarnings("unused") StartBundleContext c) { throw new IllegalArgumentException("bogus"); } @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) {} }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeStartBundle(null); }
@Test public void testFinishBundleException() throws Exception { DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @FinishBundle public void finishBundle(@SuppressWarnings("unused") FinishBundleContext c) { throw new IllegalArgumentException("bogus"); } @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) {} }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeFinishBundle(null); }
@Test public void testOnTimerHelloWord() throws Exception { final String timerId = "my-timer-id"; class SimpleTimerDoFn extends DoFn<String, String> { public String status = "not yet"; @TimerId(timerId) private final TimerSpec myTimer = TimerSpecs.timer(TimeDomain.PROCESSING_TIME); @ProcessElement public void process(ProcessContext c) {} @OnTimer(timerId) public void onMyTimer() { status = "OK now"; } } SimpleTimerDoFn fn = new SimpleTimerDoFn(); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeOnTimer(timerId, "", mockArgumentProvider); assertThat(fn.status, equalTo("OK now")); }
@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)); }