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

Example 1 with OutputReceiver

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

the class DoFnInvokersTest method testSplittableDoFnWithAllMethods.

@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());
    when(fn.getSize()).thenReturn(2.0);
    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;
        }
    }));
    assertEquals(2.0, invoker.invokeGetSize(mockArgumentProvider), 0.0001);
}

Example 2 with OutputReceiver

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

the class QueryChangeStreamActionTest method setUp.

@Before
public void setUp() throws Exception {
    changeStreamDao = mock(ChangeStreamDao.class);
    partitionMetadataDao = mock(PartitionMetadataDao.class);
    changeStreamRecordMapper = mock(ChangeStreamRecordMapper.class);
    partitionMetadataMapper = mock(PartitionMetadataMapper.class);
    dataChangeRecordAction = mock(DataChangeRecordAction.class);
    heartbeatRecordAction = mock(HeartbeatRecordAction.class);
    childPartitionsRecordAction = mock(ChildPartitionsRecordAction.class);
    action = new QueryChangeStreamAction(changeStreamDao, partitionMetadataDao, changeStreamRecordMapper, partitionMetadataMapper, dataChangeRecordAction, heartbeatRecordAction, childPartitionsRecordAction);
    final Struct row = mock(Struct.class);
    partition = PartitionMetadata.newBuilder().setPartitionToken(PARTITION_TOKEN).setParentTokens(Sets.newHashSet("parentToken")).setStartTimestamp(PARTITION_START_TIMESTAMP).setEndTimestamp(PARTITION_END_TIMESTAMP).setHeartbeatMillis(PARTITION_HEARTBEAT_MILLIS).setState(SCHEDULED).setWatermark(WATERMARK_TIMESTAMP).setScheduledAt(Timestamp.now()).build();
    restriction = mock(OffsetRange.class);
    restrictionTracker = mock(RestrictionTracker.class);
    outputReceiver = mock(OutputReceiver.class);
    watermarkEstimator = mock(ManualWatermarkEstimator.class);
    bundleFinalizer = new BundleFinalizerStub();
    when(restrictionTracker.currentRestriction()).thenReturn(restriction);
    when(restriction.getFrom()).thenReturn(10L);
    when(partitionMetadataDao.getPartition(PARTITION_TOKEN)).thenReturn(row);
    when(partitionMetadataMapper.from(row)).thenReturn(partition);
}

Example 3 with OutputReceiver

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

the class ReadChangeStreamPartitionDoFnTest method setUp.

@Before
public void setUp() {
    final DaoFactory daoFactory = mock(DaoFactory.class);
    final MapperFactory mapperFactory = mock(MapperFactory.class);
    final ChangeStreamMetrics metrics = mock(ChangeStreamMetrics.class);
    final ActionFactory actionFactory = mock(ActionFactory.class);
    final PartitionMetadataDao partitionMetadataDao = mock(PartitionMetadataDao.class);
    final ChangeStreamDao changeStreamDao = mock(ChangeStreamDao.class);
    final ChangeStreamRecordMapper changeStreamRecordMapper = mock(ChangeStreamRecordMapper.class);
    final PartitionMetadataMapper partitionMetadataMapper = mock(PartitionMetadataMapper.class);
    dataChangeRecordAction = mock(DataChangeRecordAction.class);
    heartbeatRecordAction = mock(HeartbeatRecordAction.class);
    childPartitionsRecordAction = mock(ChildPartitionsRecordAction.class);
    queryChangeStreamAction = mock(QueryChangeStreamAction.class);
    doFn = new ReadChangeStreamPartitionDoFn(daoFactory, mapperFactory, actionFactory, metrics);
    partition = PartitionMetadata.newBuilder().setPartitionToken(PARTITION_TOKEN).setParentTokens(Sets.newHashSet("parentToken")).setStartTimestamp(PARTITION_START_TIMESTAMP).setEndTimestamp(PARTITION_END_TIMESTAMP).setHeartbeatMillis(PARTITION_HEARTBEAT_MILLIS).setState(SCHEDULED).setWatermark(PARTITION_START_TIMESTAMP).setScheduledAt(Timestamp.now()).build();
    restriction = mock(OffsetRange.class);
    restrictionTracker = mock(RestrictionTracker.class);
    outputReceiver = mock(OutputReceiver.class);
    watermarkEstimator = mock(ManualWatermarkEstimator.class);
    bundleFinalizer = mock(BundleFinalizer.class);
    when(restrictionTracker.currentRestriction()).thenReturn(restriction);
    when(daoFactory.getPartitionMetadataDao()).thenReturn(partitionMetadataDao);
    when(daoFactory.getChangeStreamDao()).thenReturn(changeStreamDao);
    when(mapperFactory.changeStreamRecordMapper()).thenReturn(changeStreamRecordMapper);
    when(mapperFactory.partitionMetadataMapper()).thenReturn(partitionMetadataMapper);
    when(actionFactory.dataChangeRecordAction()).thenReturn(dataChangeRecordAction);
    when(actionFactory.heartbeatRecordAction(metrics)).thenReturn(heartbeatRecordAction);
    when(actionFactory.childPartitionsRecordAction(partitionMetadataDao, metrics)).thenReturn(childPartitionsRecordAction);
    when(actionFactory.queryChangeStreamAction(changeStreamDao, partitionMetadataDao, changeStreamRecordMapper, partitionMetadataMapper, dataChangeRecordAction, heartbeatRecordAction, childPartitionsRecordAction)).thenReturn(queryChangeStreamAction);
    doFn.setup();
}

Example 4 with OutputReceiver

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

the class DoFnInvokersTest method testSplittableDoFnWithHasDefaultMethods.

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