diff --git a/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java b/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java index 86cb7836e2..b0cafff724 100644 --- a/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java +++ b/data-prepper-plugins/kinesis-source/src/main/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisService.java @@ -41,6 +41,7 @@ import software.amazon.kinesis.exceptions.KinesisClientLibDependencyException; import software.amazon.kinesis.exceptions.ThrottlingException; import software.amazon.kinesis.processor.ShardRecordProcessorFactory; +import software.amazon.kinesis.retrieval.RetrievalConfig; import software.amazon.kinesis.retrieval.polling.PollingConfig; import java.time.Duration; @@ -185,9 +186,10 @@ public Scheduler createScheduler(final Buffer> buffer) { .tableName(tableName) .namespace(kclMetricsNamespaceName); + RetrievalConfig retrievalConfig = configsBuilder.retrievalConfig(); ConsumerStrategy consumerStrategy = kinesisSourceConfig.getConsumerStrategy(); if (consumerStrategy == ConsumerStrategy.POLLING) { - configsBuilder.retrievalConfig().retrievalSpecificConfig( + retrievalConfig = configsBuilder.retrievalConfig().retrievalSpecificConfig( new PollingConfig(kinesisClient) .maxRecords(kinesisSourceConfig.getPollingConfig().getMaxPollingRecords()) .idleTimeBetweenReadsInMillis( @@ -203,7 +205,7 @@ public Scheduler createScheduler(final Buffer> buffer) { configsBuilder.lifecycleConfig(), configsBuilder.metricsConfig(), configsBuilder.processorConfig(), - configsBuilder.retrievalConfig() + retrievalConfig ); } } diff --git a/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java b/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java index 95d01594ea..896faa4155 100644 --- a/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java +++ b/data-prepper-plugins/kinesis-source/src/test/java/org/opensearch/dataprepper/plugins/kinesis/source/KinesisServiceTest.java @@ -42,6 +42,7 @@ import software.amazon.kinesis.common.InitialPositionInStream; import software.amazon.kinesis.coordinator.Scheduler; import software.amazon.kinesis.metrics.MetricsLevel; +import software.amazon.kinesis.retrieval.polling.PollingConfig; import java.time.Duration; import java.time.Instant; @@ -276,6 +277,30 @@ void testCreateSchedulerWithPollingStrategy() { verify(workerIdentifierGenerator, times(1)).generate(); } + @Test + void testCreateSchedulerWithPollingStrategyAndPollingConfig() { + when(kinesisSourceConfig.getConsumerStrategy()).thenReturn(ConsumerStrategy.POLLING); + when(kinesisSourceConfig.getPollingConfig()).thenReturn(kinesisStreamPollingConfig); + KinesisService kinesisService = new KinesisService(kinesisSourceConfig, kinesisClientFactory, pluginMetrics, pluginFactory, + pipelineDescription, acknowledgementSetManager, kinesisLeaseConfigSupplier, workerIdentifierGenerator); + Scheduler schedulerObjectUnderTest = kinesisService.createScheduler(buffer); + + assertEquals(kinesisService.getApplicationName(), pipelineName); + assertNotNull(schedulerObjectUnderTest); + assertNotNull(schedulerObjectUnderTest.checkpointConfig()); + assertNotNull(schedulerObjectUnderTest.leaseManagementConfig()); + assertSame(schedulerObjectUnderTest.leaseManagementConfig().initialPositionInStream().getInitialPositionInStream(), InitialPositionInStream.TRIM_HORIZON); + assertNotNull(schedulerObjectUnderTest.lifecycleConfig()); + assertNotNull(schedulerObjectUnderTest.metricsConfig()); + assertSame(schedulerObjectUnderTest.metricsConfig().metricsLevel(), MetricsLevel.DETAILED); + assertNotNull(schedulerObjectUnderTest.processorConfig()); + assertNotNull(schedulerObjectUnderTest.retrievalConfig().retrievalSpecificConfig()); + assertEquals(((PollingConfig)schedulerObjectUnderTest.retrievalConfig().retrievalSpecificConfig()).maxRecords(), kinesisStreamPollingConfig.getMaxPollingRecords()); + assertEquals(((PollingConfig)schedulerObjectUnderTest.retrievalConfig().retrievalSpecificConfig()).idleTimeBetweenReadsInMillis(), kinesisStreamPollingConfig.getIdleTimeBetweenReads().toMillis()); + assertNotNull(schedulerObjectUnderTest.retrievalConfig()); + verify(workerIdentifierGenerator, times(1)).generate(); + } + @Test void testServiceStartNullBufferThrows() { KinesisService kinesisService = createObjectUnderTest();