arunlakshman commented on code in PR #198:
URL:
https://github.com/apache/flink-connector-aws/pull/198#discussion_r2045585443
##########
flink-connector-aws/flink-connector-aws-kinesis-streams/src/test/java/org/apache/flink/connector/kinesis/source/reader/KinesisStreamsSourceReaderTest.java:
##########
@@ -147,4 +151,120 @@ void testAddSplitsRegistersAndUpdatesShardMetricGroup()
throws Exception {
TestUtil.assertMillisBehindLatest(
split, TestUtil.MILLIS_BEHIND_LATEST_TEST_VALUE,
metricListener);
}
+
+ @Test
+ void testSnapshotStateWithFinishedSplits() throws Exception {
+ // Create and add a split
+ KinesisShardSplit split = getTestSplit();
+ List<KinesisShardSplit> splits = Collections.singletonList(split);
+ sourceReader.addSplits(splits);
+
+ // Set checkpoint ID by taking initial snapshot
+ List<KinesisShardSplit> initialSnapshot =
sourceReader.snapshotState(1L);
+ assertThat(initialSnapshot).hasSize(1).containsExactly(split);
+
+ // Simulate split finishing
+ Map<String, KinesisShardSplitState> finishedSplits = new HashMap<>();
+ finishedSplits.put(split.splitId(), new KinesisShardSplitState(split));
+ sourceReader.onSplitFinished(finishedSplits);
+
+ // Take another snapshot
+ List<KinesisShardSplit> snapshotSplits =
sourceReader.snapshotState(2L);
+ List<KinesisShardSplit> snapshotFinishedSplits =
+ snapshotSplits.stream()
+ .filter(KinesisShardSplit::isFinished)
+ .collect(Collectors.toList());
+ // Verify we have 2 splits - the original split and the finished split
+ assertThat(snapshotSplits).hasSize(2);
+ assertThat(snapshotFinishedSplits)
+ .hasSize(1)
+ .allSatisfy(
+ s -> {
+ assertThat(s.splitId()).isEqualTo(split.splitId());
+ });
+ }
+
+ @Test
+ void testAddSplitsWithStateRestoration() throws Exception {
+ KinesisShardSplit finishedSplit1 =
getFinishedTestSplit("finished-split-1");
+ KinesisShardSplit finishedSplit2 =
getFinishedTestSplit("finished-split-2");
+
+ // Create active split
+ KinesisShardSplit activeSplit = getTestSplit();
+
+ List<KinesisShardSplit> allSplits =
+ Arrays.asList(finishedSplit1, finishedSplit2, activeSplit);
+
+ // Clear any previous events
+ testingReaderContext.clearSentEvents();
+
+ // Add splits
+ sourceReader.addSplits(allSplits);
+
+ // Verify finished events were sent
+ List<SourceEvent> events = testingReaderContext.getSentEvents();
+ assertThat(events)
+ .hasSize(2)
+ .allMatch(e -> e instanceof SplitsFinishedEvent)
+ .allSatisfy(
+ e -> {
+ SplitsFinishedEvent event = (SplitsFinishedEvent)
e;
+ assertThat(event.getFinishedSplitIds()).hasSize(1);
+ assertThat(event.getFinishedSplitIds())
+ .containsAnyOf("finished-split-1",
"finished-split-2");
Review Comment:
Good catch. I have updated in the next commit
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscr...@flink.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
