1616
1717package org .springframework .integration .aws .kinesis ;
1818
19+ import java .util .Collections ;
1920import java .util .List ;
21+ import java .util .Map ;
22+ import java .util .concurrent .CompletableFuture ;
2023
2124import org .junit .jupiter .api .AfterAll ;
2225import org .junit .jupiter .api .BeforeAll ;
2326import org .junit .jupiter .api .Test ;
2427import software .amazon .awssdk .core .SdkBytes ;
28+ import software .amazon .awssdk .core .waiters .WaiterResponse ;
2529import software .amazon .awssdk .services .cloudwatch .CloudWatchAsyncClient ;
2630import software .amazon .awssdk .services .dynamodb .DynamoDbAsyncClient ;
31+ import software .amazon .awssdk .services .dynamodb .model .AttributeDefinition ;
32+ import software .amazon .awssdk .services .dynamodb .model .AttributeValue ;
33+ import software .amazon .awssdk .services .dynamodb .model .CreateTableRequest ;
34+ import software .amazon .awssdk .services .dynamodb .model .KeySchemaElement ;
35+ import software .amazon .awssdk .services .dynamodb .model .KeyType ;
36+ import software .amazon .awssdk .services .dynamodb .model .ProvisionedThroughput ;
37+ import software .amazon .awssdk .services .dynamodb .model .PutItemRequest ;
38+ import software .amazon .awssdk .services .dynamodb .model .PutItemResponse ;
39+ import software .amazon .awssdk .services .dynamodb .model .ScalarAttributeType ;
2740import software .amazon .awssdk .services .kinesis .KinesisAsyncClient ;
2841import software .amazon .awssdk .services .kinesis .model .Consumer ;
42+ import software .amazon .awssdk .services .kinesis .model .DescribeStreamResponse ;
2943import software .amazon .kinesis .common .InitialPositionInStream ;
3044import software .amazon .kinesis .common .InitialPositionInStreamExtended ;
3145import software .amazon .kinesis .metrics .MetricsFactory ;
3751import org .springframework .context .annotation .Configuration ;
3852import org .springframework .integration .IntegrationMessageHeaderAccessor ;
3953import org .springframework .integration .aws .LocalstackContainerTest ;
54+ import org .springframework .integration .aws .inbound .kinesis .CheckpointMode ;
4055import org .springframework .integration .aws .inbound .kinesis .KclMessageDrivenChannelAdapter ;
56+ import org .springframework .integration .aws .inbound .kinesis .ListenerMode ;
4157import org .springframework .integration .aws .support .AwsHeaders ;
4258import org .springframework .integration .channel .QueueChannel ;
4359import org .springframework .integration .config .EnableIntegration ;
5369 * @author Artem Bilan
5470 * @author Siddharth Jain
5571 * @author Minkyu Moon
56- *
5772 * @since 3.0
5873 */
5974@ SpringJUnitConfig
6075@ DirtiesContext
6176public class KclMessageDrivenChannelAdapterTests implements LocalstackContainerTest {
6277
6378 private static final String TEST_STREAM = "TestStreamKcl" ;
79+ public static final String LEASE_TABLE_NAME = "test_table" ;
80+ public static final String TEST_DATA = "test data" ;
6481
6582 private static KinesisAsyncClient AMAZON_KINESIS ;
6683
@@ -80,10 +97,10 @@ static void setup() {
8097 DYNAMO_DB = LocalstackContainerTest .dynamoDbClient ();
8198 CLOUD_WATCH = LocalstackContainerTest .cloudWatchClient ();
8299
83- AMAZON_KINESIS . createStream ( request -> request . streamName ( TEST_STREAM ). shardCount ( 1 ))
84- . thenCompose ( result ->
85- AMAZON_KINESIS . waiter (). waitUntilStreamExists ( request -> request . streamName ( TEST_STREAM )) )
86- .join ();
100+ CompletableFuture . allOf (
101+ initialiseStream ( TEST_STREAM ),
102+ initialiseLeaseTableFor ( LEASE_TABLE_NAME )
103+ ) .join ();
87104 }
88105
89106 @ AfterAll
@@ -96,20 +113,36 @@ static void tearDown() {
96113 }
97114
98115 @ Test
99- void kclChannelAdapterReceivesRecords () {
100- String testData = "test data" ;
116+ void kclChannelAdapterReceivesBatchedRecords () {
117+ this .kclMessageDrivenChannelAdapter .setListenerMode (ListenerMode .batch );
118+ this .kclMessageDrivenChannelAdapter .setCheckpointMode (CheckpointMode .batch );
119+
120+ Message <?> received = verifyRecordReceived (TEST_DATA );
121+ assertThat (received .getPayload ()).isEqualTo (Collections .singletonList (TEST_DATA ));
122+ List <?> receivedSequences = received .getHeaders ().get (AwsHeaders .RECEIVED_SEQUENCE_NUMBER , List .class );
123+ assertThat (receivedSequences ).isNotEmpty ();
124+ }
125+
126+ @ Test
127+ void kclChannelAdapterReceivesSingleRecord () {
101128
129+ this .kclMessageDrivenChannelAdapter .setListenerMode (ListenerMode .record );
130+ this .kclMessageDrivenChannelAdapter .setCheckpointMode (CheckpointMode .record );
131+
132+ Message <?> receive = verifyRecordReceived (TEST_DATA );
133+ assertThat (receive .getPayload ()).isEqualTo (TEST_DATA );
134+ assertThat (receive .getHeaders ()).containsKey (IntegrationMessageHeaderAccessor .SOURCE_DATA );
135+ assertThat (receive .getHeaders ().get (AwsHeaders .RECEIVED_SEQUENCE_NUMBER , String .class )).isNotEmpty ();
136+ }
137+
138+ private Message <?> verifyRecordReceived (String testData ) {
102139 AMAZON_KINESIS .putRecord (request ->
103140 request .streamName (TEST_STREAM )
104141 .data (SdkBytes .fromUtf8String (testData ))
105142 .partitionKey ("test" ));
106143
107- // We need so long delay because KCL has a more than a minute setup phase.
108- Message <?> receive = this .kinesisReceiveChannel .receive (300_000 );
144+ Message <?> receive = this .kinesisReceiveChannel .receive (5000 );
109145 assertThat (receive ).isNotNull ();
110- assertThat (receive .getPayload ()).isEqualTo (testData );
111- assertThat (receive .getHeaders ()).containsKey (IntegrationMessageHeaderAccessor .SOURCE_DATA );
112- assertThat (receive .getHeaders ().get (AwsHeaders .RECEIVED_SEQUENCE_NUMBER , String .class )).isNotEmpty ();
113146
114147 List <Consumer > streamConsumers =
115148 AMAZON_KINESIS .describeStream (r -> r .streamName (TEST_STREAM ))
@@ -120,10 +153,11 @@ void kclChannelAdapterReceivesRecords() {
120153 .consumers ();
121154
122155 // Because FanOut is false, there would be no Stream Consumers.
123- assertThat (streamConsumers ).hasSize ( 0 );
156+ assertThat (streamConsumers ).isEmpty ( );
124157
125158 List <String > tableNames = DYNAMO_DB .listTables ().join ().tableNames ();
126- assertThat (tableNames ).contains ("test_table" );
159+ assertThat (tableNames ).contains (LEASE_TABLE_NAME );
160+ return receive ;
127161 }
128162
129163 @ Test
@@ -176,10 +210,10 @@ public void pollingMaxRecordsIsPropagated() {
176210 public static class TestConfiguration {
177211
178212 @ Bean
179- public KclMessageDrivenChannelAdapter kclMessageDrivenChannelAdapter () {
213+ public KclMessageDrivenChannelAdapter kclMessageDrivenChannelAdapter (PollableChannel kinesisReceiveChannel ) {
180214 KclMessageDrivenChannelAdapter adapter =
181215 new KclMessageDrivenChannelAdapter (AMAZON_KINESIS , CLOUD_WATCH , DYNAMO_DB , TEST_STREAM );
182- adapter .setOutputChannel (kinesisReceiveChannel () );
216+ adapter .setOutputChannel (kinesisReceiveChannel );
183217 adapter .setStreamInitialSequence (
184218 InitialPositionInStreamExtended .newInitialPosition (InitialPositionInStream .TRIM_HORIZON ));
185219 adapter .setConverter (String ::new );
@@ -202,7 +236,45 @@ public KclMessageDrivenChannelAdapter kclMessageDrivenChannelAdapter() {
202236 public PollableChannel kinesisReceiveChannel () {
203237 return new QueueChannel ();
204238 }
239+ }
205240
241+ private static CompletableFuture <WaiterResponse <DescribeStreamResponse >> initialiseStream (String streamName ) {
242+ return AMAZON_KINESIS .createStream (request -> request .streamName (streamName ).shardCount (1 ))
243+ .thenCompose (
244+ result -> AMAZON_KINESIS .waiter ().waitUntilStreamExists (request -> request .streamName (streamName )));
206245 }
207246
247+ /**
248+ * Initialises the lease table to improve KCL initialisation time
249+ */
250+ private static CompletableFuture <PutItemResponse > initialiseLeaseTableFor (String leaseTableName ) {
251+ return DYNAMO_DB .createTable (CreateTableRequest .builder ()
252+ .tableName (leaseTableName )
253+ .attributeDefinitions (AttributeDefinition .builder ()
254+ .attributeName ("leaseKey" )
255+ .attributeType (ScalarAttributeType .S )
256+ .build ())
257+ .keySchema (KeySchemaElement .builder ()
258+ .attributeName ("leaseKey" )
259+ .keyType (KeyType .HASH )
260+ .build ())
261+ .provisionedThroughput (ProvisionedThroughput .builder ()
262+ .readCapacityUnits (1L )
263+ .writeCapacityUnits (1L )
264+ .build ())
265+ .build ())
266+ .thenCompose (
267+ result -> DYNAMO_DB .waiter ().waitUntilTableExists (request -> request .tableName (leaseTableName )))
268+ .thenCompose (describeTableResponseWaiterResponse -> DYNAMO_DB .putItem (PutItemRequest .builder ()
269+ .tableName (leaseTableName )
270+ .item (Map .of (
271+ "leaseKey" , AttributeValue .fromS ("shardId-000000000000" ),
272+ "checkpoint" , AttributeValue .fromS ("TRIM_HORIZON" ),
273+ "leaseCounter" , AttributeValue .fromN ("1" ),
274+ "startingHashKey" , AttributeValue .fromS ("0" ),
275+ "ownerSwitchesSinceCheckpoint" , AttributeValue .fromN ("0" ),
276+ "checkpointSubSequenceNumber" , AttributeValue .fromN ("0" )
277+ ))
278+ .build ()));
279+ }
208280}
0 commit comments