Examples with QueryChangeStreamAction org.apache.beam.sdk.io.gcp.spanner.changestreams.action.QueryChangeStreamAction used on opensource projects

Example 1 with QueryChangeStreamAction

use of org.apache.beam.sdk.io.gcp.spanner.changestreams.action.QueryChangeStreamAction 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();
}