use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord in project beam by apache.
the class QueryChangeStreamActionTest method testQueryChangeStreamWithDataChangeRecord.
@Test
public void testQueryChangeStreamWithDataChangeRecord() {
final Struct rowAsStruct = mock(Struct.class);
final ChangeStreamResultSetMetadata resultSetMetadata = mock(ChangeStreamResultSetMetadata.class);
final ChangeStreamResultSet resultSet = mock(ChangeStreamResultSet.class);
final DataChangeRecord record1 = mock(DataChangeRecord.class);
final DataChangeRecord record2 = mock(DataChangeRecord.class);
when(record1.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
when(record2.getRecordTimestamp()).thenReturn(PARTITION_START_TIMESTAMP);
when(changeStreamDao.changeStreamQuery(PARTITION_TOKEN, PARTITION_START_TIMESTAMP, PARTITION_END_TIMESTAMP, PARTITION_HEARTBEAT_MILLIS)).thenReturn(resultSet);
when(resultSet.next()).thenReturn(true);
when(resultSet.getCurrentRowAsStruct()).thenReturn(rowAsStruct);
when(resultSet.getMetadata()).thenReturn(resultSetMetadata);
when(changeStreamRecordMapper.toChangeStreamRecords(partition, rowAsStruct, resultSetMetadata)).thenReturn(Arrays.asList(record1, record2));
when(dataChangeRecordAction.run(partition, record1, restrictionTracker, outputReceiver, watermarkEstimator)).thenReturn(Optional.empty());
when(dataChangeRecordAction.run(partition, record2, restrictionTracker, outputReceiver, watermarkEstimator)).thenReturn(Optional.of(ProcessContinuation.stop()));
when(watermarkEstimator.currentWatermark()).thenReturn(WATERMARK);
final ProcessContinuation result = action.run(partition, restrictionTracker, outputReceiver, watermarkEstimator, bundleFinalizer);
assertEquals(ProcessContinuation.stop(), result);
verify(dataChangeRecordAction).run(partition, record1, restrictionTracker, outputReceiver, watermarkEstimator);
verify(dataChangeRecordAction).run(partition, record2, restrictionTracker, outputReceiver, watermarkEstimator);
verify(partitionMetadataDao).updateWatermark(PARTITION_TOKEN, WATERMARK_TIMESTAMP);
verify(heartbeatRecordAction, never()).run(any(), any(), any(), any());
verify(childPartitionsRecordAction, never()).run(any(), any(), any(), any());
verify(restrictionTracker, never()).tryClaim(any());
}
use of org.apache.beam.sdk.io.gcp.spanner.changestreams.model.DataChangeRecord in project beam by apache.
the class DataChangeRecordActionTest method testRestrictionClaimed.
@Test
public void testRestrictionClaimed() {
final String partitionToken = "partitionToken";
final Timestamp timestamp = Timestamp.ofTimeMicroseconds(10L);
final Instant instant = new Instant(timestamp.toSqlTimestamp().getTime());
final DataChangeRecord record = mock(DataChangeRecord.class);
when(record.getCommitTimestamp()).thenReturn(timestamp);
when(tracker.tryClaim(10L)).thenReturn(true);
when(partition.getPartitionToken()).thenReturn(partitionToken);
final Optional<ProcessContinuation> maybeContinuation = action.run(partition, record, tracker, outputReceiver, watermarkEstimator);
assertEquals(Optional.empty(), maybeContinuation);
verify(outputReceiver).outputWithTimestamp(record, instant);
verify(watermarkEstimator).setWatermark(instant);
}
Aggregations