Kinesis scheduler test

.gitignore

@@ -29,3 +29,5 @@ hs_err_pid*
 .classpath
 .settings
 .project
+
+.vscode

pom.xml

@@ -1,5 +1,4 @@
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-    xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
     <modelVersion>4.0.0</modelVersion>
 
     <groupId>cloud.localstack</groupId>
@@ -207,6 +206,30 @@
             <scope>provided</scope>
         </dependency>
 
+        <dependency>
+            <groupId>software.amazon.kinesis</groupId>
+            <artifactId>amazon-kinesis-client</artifactId>
+            <version>2.2.9</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-core</artifactId>
+            <version>2.14.1</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-api</artifactId>
+            <version>2.14.1</version>
+            <scope>provided</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.logging.log4j</groupId>
+            <artifactId>log4j-slf4j-impl</artifactId>
+            <version>2.14.1</version>
+            <scope>provided</scope>
+        </dependency>
 
         <!-- test dependencies -->
         <dependency>
@@ -287,8 +310,7 @@
         </profile>
         <profile>
             <id>awssdkv1</id>
-            <dependencies>
-            </dependencies>
+            <dependencies></dependencies>
         </profile>
         <profile>
             <id>awssdkv2</id>
@@ -308,6 +330,11 @@
                     <artifactId>kinesis</artifactId>
                     <version>${aws.sdkv2.version}</version>
                 </dependency>
+                <dependency>
+                    <groupId>software.amazon.kinesis</groupId>
+                    <artifactId>amazon-kinesis-client</artifactId>
+                    <version>2.2.9</version>
+                </dependency>
             </dependencies>
         </profile>
     </profiles>

src/test/java/cloud/localstack/awssdkv2/KinesisSchedulerTest.java

@@ -0,0 +1,94 @@
+package cloud.localstack.awssdkv2;
+
+import cloud.localstack.awssdkv2.consumer.DeliveryStatusRecordProcessorFactory;
+import cloud.localstack.awssdkv2.consumer.EventProcessor;
+import cloud.localstack.docker.annotation.LocalstackDockerProperties;
+import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.core.SdkSystemSetting;
+import software.amazon.awssdk.services.cloudwatch.CloudWatchAsyncClient;
+import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
+import software.amazon.awssdk.services.kinesis.KinesisAsyncClient;
+import software.amazon.awssdk.services.kinesis.model.CreateStreamRequest;
+import software.amazon.awssdk.services.kinesis.model.CreateStreamResponse;
+import software.amazon.awssdk.services.kinesis.model.PutRecordRequest;
+import software.amazon.awssdk.services.kinesis.model.PutRecordResponse;
+import software.amazon.kinesis.common.ConfigsBuilder;
+import software.amazon.kinesis.coordinator.Scheduler;
+import software.amazon.kinesis.metrics.NullMetricsFactory;
+
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.*;
+import java.util.concurrent.TimeUnit;
+
+@LocalstackDockerProperties(ignoreDockerRunErrors = true)
+public class KinesisSchedulerTest extends PowerMockLocalStack {
+  String streamName = "test" + UUID.randomUUID().toString();
+  String workerId = UUID.randomUUID().toString();
+  String testMessage = "hello, world";
+  Integer consumerCreationTime = 15; //35 for aws
+
+  @Before
+  public void mockServicesForScheduler() {
+    // System.setProperty(SdkSystemSetting.CBOR_ENABLED.property(), "false");
+    PowerMockLocalStack.mockCloudWatchAsyncClient();
+    PowerMockLocalStack.mockDynamoDBAsync();
+    PowerMockLocalStack.mockKinesisAsync();
+  }
+
+  @Test
+  public void schedulerTest() throws Exception {
+
+    KinesisAsyncClient kinesisAsyncClient = KinesisAsyncClient.create();
+    DynamoDbAsyncClient dynamoAsyncClient = DynamoDbAsyncClient.create();
+    CloudWatchAsyncClient cloudWatchAsyncClient = CloudWatchAsyncClient.create();
+
+    createStream(kinesisAsyncClient);
+    TimeUnit.SECONDS.sleep(2);
+
+    EventProcessor eventProcessor = new EventProcessor();
+    DeliveryStatusRecordProcessorFactory processorFactory = new DeliveryStatusRecordProcessorFactory(eventProcessor);
+
+    ConfigsBuilder configsBuilder = new ConfigsBuilder(streamName, streamName, kinesisAsyncClient, dynamoAsyncClient,
+        cloudWatchAsyncClient, workerId, processorFactory);
+    Scheduler scheduler = createScheduler(configsBuilder);
+
+    new Thread(scheduler).start();
+    TimeUnit.SECONDS.sleep(consumerCreationTime);
+
+    putRecord(kinesisAsyncClient);
+    TimeUnit.SECONDS.sleep(5);
+
+    scheduler.shutdown();
+    Assert.assertTrue(eventProcessor.CONSUMER_CREATED);
+    Assert.assertTrue(eventProcessor.RECORD_RECEIVED);
+    Assert.assertTrue(eventProcessor.messages.size() > 0);
+    Assert.assertEquals(eventProcessor.messages.get(0), testMessage);
+  }
+
+  public Scheduler createScheduler(ConfigsBuilder configsBuilder) {
+    return new Scheduler(configsBuilder.checkpointConfig(), configsBuilder.coordinatorConfig(),
+        configsBuilder.leaseManagementConfig(), configsBuilder.lifecycleConfig(),
+        configsBuilder.metricsConfig().metricsFactory(new NullMetricsFactory()), configsBuilder.processorConfig(),
+        configsBuilder.retrievalConfig());
+  }
+
+  public void createStream(KinesisAsyncClient kinesisClient) throws Exception {
+    CreateStreamRequest request = CreateStreamRequest.builder().streamName(streamName).shardCount(1).build();
+    CreateStreamResponse response = kinesisClient.createStream(request).get();
+
+    Assert.assertNotNull(response);
+  }
+
+  public void putRecord(KinesisAsyncClient kinesisClient) throws Exception {
+    System.out.println("PUTTING RECORD");
+    PutRecordRequest request = PutRecordRequest.builder().partitionKey("partitionkey").streamName(streamName)
+        .data(SdkBytes.fromUtf8String(testMessage)).build();
+    PutRecordResponse response = kinesisClient.putRecord(request).get();
+
+    Assert.assertNotNull(response);
+  }
+
+}

src/test/java/cloud/localstack/awssdkv2/KinesisV2ConsumerTest.java

@@ -1,6 +1,7 @@
 package cloud.localstack.awssdkv2;
 
 import cloud.localstack.LocalstackTestRunner;
+import cloud.localstack.docker.annotation.LocalstackDockerProperties;
 
 import org.junit.Assert;
 import org.junit.Test;
@@ -16,6 +17,7 @@
 import java.util.stream.Collectors;
 
 @RunWith(LocalstackTestRunner.class)
+@LocalstackDockerProperties(ignoreDockerRunErrors = true)
 public class KinesisV2ConsumerTest {
 
   @Test
@@ -49,6 +51,6 @@ public void testGetRecordCBOR() throws Exception {
   public void testGetRecordJSON() throws Exception {
     System.setProperty(SdkSystemSetting.CBOR_ENABLED.property(), "false");
     this.testGetRecordCBOR();
-    System.setProperty(SdkSystemSetting.CBOR_ENABLED.property(), "false");
+    System.setProperty(SdkSystemSetting.CBOR_ENABLED.property(), "true");
   }
 }

src/test/java/cloud/localstack/awssdkv2/consumer/DeliveryStatusProcessor.java

@@ -0,0 +1,58 @@
+package cloud.localstack.awssdkv2.consumer;
+
+import java.io.IOException;
+import java.util.logging.Logger;
+
+import software.amazon.kinesis.lifecycle.events.InitializationInput;
+import software.amazon.kinesis.lifecycle.events.LeaseLostInput;
+import software.amazon.kinesis.lifecycle.events.ProcessRecordsInput;
+import software.amazon.kinesis.lifecycle.events.ShardEndedInput;
+import software.amazon.kinesis.lifecycle.events.ShutdownRequestedInput;
+import software.amazon.kinesis.processor.ShardRecordProcessor;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+public class DeliveryStatusProcessor implements ShardRecordProcessor {
+    EventProcessor eventProcessor;
+    private static final Logger LOG = Logger.getLogger(DeliveryStatusProcessor.class.getName());
+
+    public DeliveryStatusProcessor(EventProcessor eventProcessor) {
+        this.eventProcessor = eventProcessor;
+    }
+
+    public void initialize(InitializationInput initializationInput) {
+        this.eventProcessor.CONSUMER_CREATED = true;
+    }
+
+    public void processRecords(ProcessRecordsInput processRecordsInput) {
+        LOG.info("RECORDS PROCESSING");
+        this.eventProcessor.RECORD_RECEIVED = true;
+        processRecordsInput.records().forEach(record -> {
+            try {
+                processRecord(record);
+            } catch (IOException e) {
+                e.printStackTrace();
+            }
+        });
+    }
+
+    public void processRecord(KinesisClientRecord record) throws IOException {
+        LOG.info("RECORD PROCESSING");
+        this.eventProcessor.RECORD_RECEIVED = true;
+        byte[] message = new byte[record.data().remaining()];
+        record.data().get(message);
+        String string = new String(message);
+        eventProcessor.messages.add(string);
+    }
+
+    public void processAndPublishRecord(byte[] messageStatus) throws IOException {
+    }
+
+    public void leaseLost(LeaseLostInput leaseLostInput) {
+    }
+
+    public void shardEnded(ShardEndedInput shardEndedInput) {
+    }
+
+    public void shutdownRequested(ShutdownRequestedInput shutdownRequestedInput) {
+    }
+}

src/test/java/cloud/localstack/awssdkv2/consumer/DeliveryStatusRecordProcessorFactory.java

@@ -0,0 +1,16 @@
+package cloud.localstack.awssdkv2.consumer;
+
+import software.amazon.kinesis.processor.ShardRecordProcessor;
+import software.amazon.kinesis.processor.ShardRecordProcessorFactory;
+
+public class DeliveryStatusRecordProcessorFactory implements ShardRecordProcessorFactory{
+    private final EventProcessor eventProcessor;
+
+    public DeliveryStatusRecordProcessorFactory(EventProcessor eventProcessor) {
+      this.eventProcessor = eventProcessor;
+    }
+
+    public ShardRecordProcessor shardRecordProcessor() {
+      return new DeliveryStatusProcessor(eventProcessor);
+    }
+}

src/test/java/cloud/localstack/awssdkv2/consumer/EventProcessor.java

@@ -0,0 +1,10 @@
+package cloud.localstack.awssdkv2.consumer;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class EventProcessor {
+    public Boolean CONSUMER_CREATED = false;
+    public Boolean RECORD_RECEIVED = false;
+    public List<String> messages = new ArrayList<>();
+}

src/test/resources/log4j2.xml

@@ -1,14 +1,14 @@
-<Configuration status="DEBUG">
+<Configuration status="INFO">
   <Appenders>
     <Console name="ConsoleAppender" target="SYSTEM_OUT">
       <PatternLayout pattern="%d{YYYY-MM-dd HH:mm:ss} [%t] %-5p %c:%L - %m%n" />
     </Console>
   </Appenders>
 
   <Loggers>
-    <Root level="DEBUG">
+    <Root level="INFO">
       <AppenderRef ref="ConsoleAppender"/>
     </Root>
-    <Logger name="software.amazon.awssdk" level="DEBUG" />
+    <Logger name="software.amazon.awssdk" level="INFO" />
   </Loggers>
 </Configuration>