Add paging functionality to getLogGroupNames

src/main/java/org/graylog/integrations/IntegrationsModule.java

@@ -23,6 +23,9 @@
 import org.graylog2.plugin.PluginModule;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClientBuilder;
 import software.amazon.awssdk.services.kinesis.KinesisClient;
 import software.amazon.awssdk.services.kinesis.KinesisClientBuilder;
 
@@ -71,6 +74,8 @@ protected void configure() {
         addCodec(PaloAltoCodec.NAME, PaloAltoCodec.class);
 
         addRestResource(AWSResource.class);
+
+        bind(CloudWatchLogsClientBuilder.class).toProvider(CloudWatchLogsClient::builder);
         bind(KinesisClientBuilder.class).toProvider(KinesisClient::builder);
     }
 }

src/main/java/org/graylog/integrations/aws/CloudWatchService.java

@@ -2,82 +2,41 @@
 
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClientBuilder;
+import software.amazon.awssdk.services.cloudwatchlogs.model.DescribeLogGroupsRequest;
 import software.amazon.awssdk.services.cloudwatchlogs.model.DescribeLogGroupsResponse;
-import software.amazon.awssdk.services.cloudwatchlogs.model.DescribeLogStreamsRequest;
-import software.amazon.awssdk.services.cloudwatchlogs.model.GetLogEventsRequest;
+import software.amazon.awssdk.services.cloudwatchlogs.paginators.DescribeLogGroupsIterable;
 
 import javax.inject.Inject;
 import java.util.ArrayList;
-import java.util.Iterator;
-import java.util.List;
 
 public class CloudWatchService {
 
-    private CloudWatchLogsClient logsClient;
+    private CloudWatchLogsClientBuilder logsClientBuilder;
 
     @Inject
-    public CloudWatchService() {
-
-    }
-
-    static CloudWatchLogsClient createCloudWatchLogClient(String region) {
-
-        CloudWatchLogsClient logsClient = CloudWatchLogsClient.builder()
-                //.credentialsProvider(StaticCredentialsProvider.create(basicCredentials))
-                .region(Region.of(region))
-                .build();
-        return logsClient;
-    }
-
-    public CloudWatchLogsClient createGetLogRequest(String logGroupName, String logStreamName, boolean fromStart) {
-        logsClient.getLogEvents(createGetLogEventRequest(logGroupName, logStreamName, fromStart));
-        return logsClient;
-    }
-
-    static GetLogEventsRequest createGetLogEventRequest(String logGroupName, String logStreamName, boolean fromStart) {
-        GetLogEventsRequest getLogEventsRequest = GetLogEventsRequest.builder()
-                .logGroupName(logGroupName)
-                .logStreamName(logStreamName)
-                //.startTime()
-                //.endTime()
-                //.nextToken(nextToken)
-                //.limit(logLimit)
-                .startFromHead(fromStart)
-                .build();
-        return getLogEventsRequest;
+    public CloudWatchService(CloudWatchLogsClientBuilder logsClientBuilder) {
+        this.logsClientBuilder = logsClientBuilder;
     }
 
+    /**
+     * Returns a list of log groups that exist in CloudWatch.
+     *
+     * @param region The AWS region
+     * @return A list of log groups in alphabetical order.
+     */
     public ArrayList<String> getLogGroupNames(String region) {
 
-        ArrayList<String> logGroupNames = new ArrayList<>();
-        // TODO optimize this
-        Iterator<DescribeLogGroupsResponse> logGroupsIterator = CloudWatchService.createCloudWatchLogClient(region).describeLogGroupsPaginator().iterator();
-        DescribeLogGroupsResponse response = logGroupsIterator.next();
-        for (int c = 0; c < response.logGroups().size(); c++) {
-            response.logGroups().get(c).logGroupName();
-            logGroupNames.add(response.logGroups().get(c).logGroupName());
-        }
-        return logGroupNames;
-    }
+        final CloudWatchLogsClient cloudWatchLogsClient = logsClientBuilder.region(Region.of(region)).build();
+        final DescribeLogGroupsRequest describeLogGroupsRequest = DescribeLogGroupsRequest.builder().build();
+        final DescribeLogGroupsIterable responses = cloudWatchLogsClient.describeLogGroupsPaginator(describeLogGroupsRequest);
 
-    static ArrayList<String> getStreamNameList(CloudWatchLogsClient cloudWatchLogsClient, String logGroupName) {
-        DescribeLogStreamsRequest logStreamsRequest = DescribeLogStreamsRequest.builder()
-                .logGroupName(logGroupName)
-                .build();
-        int logStreamListSize = cloudWatchLogsClient.describeLogStreams(((logStreamsRequest))).logStreams().size();
-        ArrayList<String> streamNameList = new ArrayList<>();
-        for (int c = 0; c < logStreamListSize; c++) {
-            String logStreamName = cloudWatchLogsClient.describeLogStreams(((logStreamsRequest))).logStreams().get(c).logStreamName();
-            streamNameList.add(logStreamName);
+        final ArrayList<String> groupNameList = new ArrayList<>();
+        for (DescribeLogGroupsResponse response : responses) {
+            for (int c = 0; c < response.logGroups().size(); c++) {
+                groupNameList.add(response.logGroups().get(c).logGroupName());
+            }
         }
-        return streamNameList;
-    }
-
-    List<String> fakeLogGroups() {
-
-        ArrayList<String> logGroups = new ArrayList<>();
-        logGroups.add("test-group1");
-        logGroups.add("test-group2");
-        return logGroups;
+        return groupNameList;
     }
 }

src/test/java/org.graylog.integrations/aws/CloudWatchServiceTest.java

@@ -1,35 +1,87 @@
 package org.graylog.integrations.aws;
 
+import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClient;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClientBuilder;
+import software.amazon.awssdk.services.cloudwatchlogs.model.DescribeLogGroupsRequest;
+import software.amazon.awssdk.services.cloudwatchlogs.model.DescribeLogGroupsResponse;
+import software.amazon.awssdk.services.cloudwatchlogs.model.LogGroup;
+import software.amazon.awssdk.services.cloudwatchlogs.paginators.DescribeLogGroupsIterable;
 
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
-import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.isA;
+import static org.mockito.Mockito.when;
 
 public class CloudWatchServiceTest {
 
+    @Rule
+    public MockitoRule mockitoRule = MockitoJUnit.rule();
+
+    @Mock
+    private CloudWatchLogsClientBuilder logsClientBuilder;
+
+    @Mock
+    private CloudWatchLogsClient cloudWatchLogsClient;
+
+    @Mock
+    DescribeLogGroupsIterable logGroupsIterable;
+
     private CloudWatchService cloudWatchService;
 
     @Before
     public void setUp() {
 
-        cloudWatchService = new CloudWatchService();
+        cloudWatchService = new CloudWatchService(logsClientBuilder);
     }
 
     @Test
     public void testLogGroupNames() {
 
-        List<String> logGroups = cloudWatchService.fakeLogGroups();
+        // Perform test setup. Return the builder and client when appropriate.
+        when(logsClientBuilder.region(isA(Region.class))).thenReturn(logsClientBuilder);
+        when(logsClientBuilder.build()).thenReturn(cloudWatchLogsClient);
+
+        // Create a fake response that contains three log groups.
+        DescribeLogGroupsResponse fakeLogGroupResponse = DescribeLogGroupsResponse
+                .builder()
+                .logGroups(LogGroup.builder().logGroupName("group-1").build(),
+                           LogGroup.builder().logGroupName("group-2").build(),
+                           LogGroup.builder().logGroupName("group-3").build())
+                .build();
 
-        boolean foundLogGroupName = false;
-        for (String logGroup : logGroups) {
-            if (logGroup.equals("test-group1")) {
-                foundLogGroupName = true;
+        // Mock out the response. When CloudWatchLogsClient.describeLogGroupsPaginator() is called,
+        // return two responses with six messages total.
+        List<DescribeLogGroupsResponse> responses = Arrays.asList(fakeLogGroupResponse, fakeLogGroupResponse);
+        when(logGroupsIterable.iterator()).thenReturn(responses.iterator());
+        when(cloudWatchLogsClient.describeLogGroupsPaginator(isA(DescribeLogGroupsRequest.class))).thenReturn(logGroupsIterable);
+
+        ArrayList<String> logGroupNames = cloudWatchService.getLogGroupNames("us-east-1");
+
+        // Inspect the log groups returned and verify the contents and size.
+        Assert.assertEquals("The number of groups should be because the two responses " +
+                                    "with 3 groups each were provided.", 6, logGroupNames.size());
+
+        // Loop example to verify presence of a specific log group.
+        boolean foundGroup = false;
+        for (String logGroupName : logGroupNames) {
+            if (logGroupName.equals("group-1")) {
+                foundGroup = true;
             }
         }
-        assertTrue(foundLogGroupName);
-        assertEquals(2, logGroups.size());
+        assertTrue(foundGroup);
+
+        // One line with stream.
+        assertTrue(logGroupNames.stream().anyMatch(logGroupName -> logGroupName.equals("group-2")));
     }
 }

src/test/java/org.graylog.integrations/aws/KinesisServiceTest.java

@@ -29,6 +29,7 @@ public class KinesisServiceTest {
 
     @Rule
     public MockitoRule mockitoRule = MockitoJUnit.rule();
+
     @Mock
     private KinesisClientBuilder kinesisClientBuilder;
 

src/test/java/org.graylog.integrations/aws/resources/AWSResourceTest.java

@@ -9,6 +9,7 @@
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.cloudwatchlogs.CloudWatchLogsClientBuilder;
 import software.amazon.awssdk.services.kinesis.KinesisClient;
 import software.amazon.awssdk.services.kinesis.KinesisClientBuilder;
 
@@ -24,6 +25,9 @@ public class AWSResourceTest {
     @Mock
     private KinesisClientBuilder kinesisClientBuilder;
 
+    @Mock
+    private CloudWatchLogsClientBuilder logsClientBuilder;
+
     @Mock
     KinesisClient kinesisClient;
 
@@ -35,6 +39,6 @@ public void setUp() {
         when(kinesisClientBuilder.build()).thenReturn(kinesisClient);
 
         // Set up the chain of mocks.
-        awsResource = new AWSResource(new AWSService(), new KinesisService(kinesisClientBuilder), new CloudWatchService());
+        awsResource = new AWSResource(new AWSService(), new KinesisService(kinesisClientBuilder), new CloudWatchService(logsClientBuilder));
     }
 }