Examples with DescribeLogStreamsRequest com.amazonaws.services.logs.model.DescribeLogStreamsRequest used on opensource projects

Example 11 with DescribeLogStreamsRequest

use of com.amazonaws.services.logs.model.DescribeLogStreamsRequest in project aws-athena-query-federation by awslabs.

the class CloudwatchMetadataHandlerTest method doListTables.

@Test
public void doListTables() throws TimeoutException {
    logger.info("doListTables - enter");
    when(mockAwsLogs.describeLogStreams(any(DescribeLogStreamsRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
        DescribeLogStreamsRequest request = (DescribeLogStreamsRequest) invocationOnMock.getArguments()[0];
        DescribeLogStreamsResult result = new DescribeLogStreamsResult();
        Integer nextToken;
        if (request.getNextToken() == null) {
            nextToken = 1;
        } else if (Integer.valueOf(request.getNextToken()) < 3) {
            nextToken = Integer.valueOf(request.getNextToken()) + 1;
        } else {
            nextToken = null;
        }
        List<LogStream> logStreams = new ArrayList<>();
        if (request.getNextToken() == null || Integer.valueOf(request.getNextToken()) < 3) {
            for (int i = 0; i < 10; i++) {
                LogStream nextLogStream = new LogStream();
                nextLogStream.setLogStreamName("table-" + String.valueOf(i));
                logStreams.add(nextLogStream);
            }
        }
        result.withLogStreams(logStreams);
        if (nextToken != null) {
            result.setNextToken(String.valueOf(nextToken));
        }
        return result;
    });
    ListTablesRequest req = new ListTablesRequest(identity, "queryId", "default", "schema-1", null, UNLIMITED_PAGE_SIZE_VALUE);
    ListTablesResponse res = handler.doListTables(allocator, req);
    logger.info("doListTables - {}", res.getTables());
    assertTrue(res.getTables().contains(new TableName("schema-1", "all_log_streams")));
    assertTrue(res.getTables().size() == 31);
    verify(mockAwsLogs, times(4)).describeLogStreams(any(DescribeLogStreamsRequest.class));
    verify(mockAwsLogs, times(1)).describeLogGroups(any(DescribeLogGroupsRequest.class));
    verifyNoMoreInteractions(mockAwsLogs);
    logger.info("doListTables - exit");
}