comment

Signed-off-by: Yupeng Fu <[email protected]>

Author: yupeng9

plugins/ingestion-kinesis/src/internalClusterTest/java/org/opensearch/plugin/kinesis/IngestFromKinesisIT.java

@@ -28,6 +28,14 @@
 import java.util.stream.Stream;
 
 import org.testcontainers.containers.localstack.LocalStackContainer;
+import software.amazon.awssdk.services.kinesis.model.DescribeStreamRequest;
+import software.amazon.awssdk.services.kinesis.model.DescribeStreamResponse;
+import software.amazon.awssdk.services.kinesis.model.GetRecordsRequest;
+import software.amazon.awssdk.services.kinesis.model.GetRecordsResponse;
+import software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest;
+import software.amazon.awssdk.services.kinesis.model.GetShardIteratorResponse;
+import software.amazon.awssdk.services.kinesis.model.Record;
+import software.amazon.awssdk.services.kinesis.model.ShardIteratorType;
 
 import static org.hamcrest.Matchers.is;
 import static org.awaitility.Awaitility.await;
@@ -98,7 +106,11 @@ public void testKinesisIngestion_RewindByOffset() throws InterruptedException {
         String sequenceNumber = produceData("3", "name3", "20");
         logger.info("Produced message with sequence number: {}", sequenceNumber);
         produceData("4", "name4", "21");
-        Thread.sleep(2000);
+
+        await()
+            .atMost(5, TimeUnit.SECONDS)
+            .until(() -> isRewinded(sequenceNumber));
+
 
         // create an index with ingestion source from kinesis
         createIndex(
@@ -128,4 +140,32 @@ public void testKinesisIngestion_RewindByOffset() throws InterruptedException {
             assertThat(response.getHits().getTotalHits().value(), is(2L));
         });
     }
+
+    private boolean isRewinded(String sequenceNumber) {
+        DescribeStreamResponse describeStreamResponse =
+            kinesisClient.describeStream(DescribeStreamRequest.builder().streamName(streamName).build());
+
+        String shardId = describeStreamResponse.streamDescription().shards().get(0).shardId();
+
+        GetShardIteratorRequest iteratorRequest = GetShardIteratorRequest.builder()
+            .streamName(streamName)
+            .shardId(shardId)
+            .shardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
+            .startingSequenceNumber(sequenceNumber)
+            .build();
+
+        GetShardIteratorResponse iteratorResponse = kinesisClient.getShardIterator(iteratorRequest);
+        String shardIterator = iteratorResponse.shardIterator();
+
+        // Use the iterator to read the record
+        GetRecordsRequest recordsRequest = GetRecordsRequest.builder().shardIterator(shardIterator).limit(1)  // Adjust as needed
+            .build();
+
+        GetRecordsResponse recordsResponse = kinesisClient.getRecords(recordsRequest);
+        List<Record> records = recordsResponse.records();
+        if (records.size() != 1) {
+            return false;
+        }
+        return records.get(0).partitionKey().equals("3");
+    }
 }