Modify the exception handling to be more descriptive when stream is n…

…ot available. Signed-off-by: Souvik Bose <souvbose@amazon.com>
opensearch-project · Feb 24, 2025 · b1af3c4 · b1af3c4
1 parent 94d1f15
commit b1af3c4
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 1 deletion.
diff --git a/...nsearch/dataprepper/plugins/kinesis/source/exceptions/KinesisStreamNotFoundException.java b/...nsearch/dataprepper/plugins/kinesis/source/exceptions/KinesisStreamNotFoundException.java
@@ -0,0 +1,7 @@
+package org.opensearch.dataprepper.plugins.kinesis.source.exceptions;
+
+public class KinesisStreamNotFoundException extends RuntimeException {
+    public KinesisStreamNotFoundException(final String errorMessage) {
+        super(errorMessage);
+    }
+}
diff --git a/...a/org/opensearch/dataprepper/plugins/kinesis/source/processor/KinesisRecordProcessor.java b/...a/org/opensearch/dataprepper/plugins/kinesis/source/processor/KinesisRecordProcessor.java
@@ -23,6 +23,7 @@
 import org.opensearch.dataprepper.plugins.kinesis.source.configuration.KinesisSourceConfig;
 import org.opensearch.dataprepper.plugins.kinesis.source.configuration.KinesisStreamConfig;
 import org.opensearch.dataprepper.plugins.kinesis.source.converter.KinesisRecordConverter;
+import org.opensearch.dataprepper.plugins.kinesis.source.exceptions.KinesisStreamNotFoundException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.kinesis.common.StreamIdentifier;
@@ -111,7 +112,11 @@ public KinesisRecordProcessor(final BufferAccumulator<Record<Event>> bufferAccum
     }
 
     private KinesisStreamConfig getStreamConfig(final KinesisSourceConfig kinesisSourceConfig) {
-        return kinesisSourceConfig.getStreams().stream().filter(streamConfig -> streamConfig.getName().equals(streamIdentifier.streamName())).findAny().get();
+        final Optional<KinesisStreamConfig> kinesisStreamConfig = kinesisSourceConfig.getStreams().stream().filter(streamConfig -> streamConfig.getName().equals(streamIdentifier.streamName())).findAny();
+        if (kinesisStreamConfig.isEmpty()) {
+            throw new KinesisStreamNotFoundException(String.format("Kinesis stream not found for %s", streamIdentifier.streamName()));
+        }
+        return kinesisStreamConfig.get();
     }
 
     @Override

diff --git a/...g/opensearch/dataprepper/plugins/kinesis/source/processor/KinesisRecordProcessorTest.java b/...g/opensearch/dataprepper/plugins/kinesis/source/processor/KinesisRecordProcessorTest.java
@@ -36,6 +36,7 @@
 import org.opensearch.dataprepper.plugins.kinesis.source.configuration.KinesisStreamConfig;
 import org.opensearch.dataprepper.plugins.kinesis.source.converter.KinesisRecordConverter;
 import org.opensearch.dataprepper.plugins.kinesis.source.converter.MetadataKeyAttributes;
+import org.opensearch.dataprepper.plugins.kinesis.source.exceptions.KinesisStreamNotFoundException;
 import software.amazon.kinesis.common.StreamIdentifier;
 import software.amazon.kinesis.exceptions.InvalidStateException;
 import software.amazon.kinesis.exceptions.ShutdownException;
@@ -53,6 +54,7 @@
 import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
@@ -62,6 +64,7 @@
 
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyDouble;
 import static org.mockito.ArgumentMatchers.eq;
@@ -617,6 +620,47 @@ void testShutdownRequestedCheckpointerThrowsNoThrowExceptionRegularCheckpoint(fi
         verify(checkpointFailures, times(1)).increment();
     }
 
+    @Test
+    void testGetStreamConfig_StreamFound() {
+        final String streamName = UUID.randomUUID().toString();
+        when(kinesisStreamConfig.getName()).thenReturn(streamName);
+        when(streamIdentifier.streamName()).thenReturn(streamName);
+        when(kinesisSourceConfig.getStreams()).thenReturn(Arrays.asList(kinesisStreamConfig));
+
+        assertDoesNotThrow(() -> new KinesisRecordProcessor(bufferAccumulator, kinesisSourceConfig,
+                acknowledgementSetManager, pluginMetrics, kinesisRecordConverter, kinesisCheckpointerTracker, streamIdentifier));
+    }
+
+    @Test
+    void testGetStreamConfig_StreamNotFound() {
+        // Arrange
+        when(streamIdentifier.streamName()).thenReturn(UUID.randomUUID().toString());
+        when(kinesisSourceConfig.getStreams()).thenReturn(Collections.emptyList());
+
+        // Act & Assert
+        assertThrows(KinesisStreamNotFoundException.class, () -> new KinesisRecordProcessor(bufferAccumulator, kinesisSourceConfig,
+                acknowledgementSetManager, pluginMetrics, kinesisRecordConverter, kinesisCheckpointerTracker, streamIdentifier));
+    }
+
+    @Test
+    void testGetStreamConfig_MultipleStreams() {
+        // Arrange
+        List<KinesisStreamConfig> streamConfigs = new ArrayList<>();
+        final String streamName1 = UUID.randomUUID().toString();
+        when(streamIdentifier.streamName()).thenReturn(streamName1);
+        KinesisStreamConfig streamConfig1 = mock(KinesisStreamConfig.class);
+        when(streamConfig1.getName()).thenReturn(streamName1);
+        streamConfigs.add(streamConfig1);
+        KinesisStreamConfig streamConfig2 = mock(KinesisStreamConfig.class);
+        when(streamConfig2.getName()).thenReturn(UUID.randomUUID().toString());
+        streamConfigs.add(streamConfig2);
+
+        when(kinesisSourceConfig.getStreams()).thenReturn(streamConfigs);
+
+        assertDoesNotThrow(() -> new KinesisRecordProcessor(bufferAccumulator, kinesisSourceConfig,
+                acknowledgementSetManager, pluginMetrics, kinesisRecordConverter, kinesisCheckpointerTracker, streamIdentifier));
+    }
+
     private List<KinesisClientRecord> createInputKinesisClientRecords() {
         List<KinesisClientRecord> kinesisClientRecords = new ArrayList<>();
         for (int i = 0; i< KinesisRecordProcessorTest.NUMBER_OF_RECORDS_TO_ACCUMULATE; i++) {