Examples with DoFn org.apache.beam.sdk.transforms.DoFn used on opensource projects

Example 71 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.

the class DoFnInvokersTest method testTruncateFnWithHasDefaultMethodsWhenUnbounded.

@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));
}

Example 72 with DoFn

use of org.apache.beam.sdk.transforms.DoFn 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(TimerDeclaration.PREFIX + timerId, "", mockArgumentProvider);
    assertThat(fn.window, equalTo(testWindow));
}

Example 73 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.

the class DoFnInvokersTest method testDefaultWatermarkEstimatorStateAndCoder.

@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<>()));
}

Example 74 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.

the class DoFnInvokersTest method testBundleFinalizer.

@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));
}

Example 75 with DoFn

use of org.apache.beam.sdk.transforms.DoFn in project beam by apache.

the class DoFnSignaturesProcessElementTest method testBadGenericWildCards.

@Test
public void testBadGenericWildCards() throws Exception {
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("DoFn<Integer, ? super Integer>.ProcessContext");
    thrown.expectMessage("must have type");
    thrown.expectMessage("DoFn<Integer, String>.ProcessContext");
    analyzeProcessElementMethod(new AnonymousMethod() {

        private void method(DoFn<Integer, ? super Integer>.ProcessContext c) {
        }
    });
}