Search in sources :

Example 11 with LogStream

use of com.amazonaws.services.logs.model.LogStream 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");
}
Also used : TableName(com.amazonaws.athena.connector.lambda.domain.TableName) InvocationOnMock(org.mockito.invocation.InvocationOnMock) ArrayList(java.util.ArrayList) ListTablesRequest(com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest) DescribeLogStreamsRequest(com.amazonaws.services.logs.model.DescribeLogStreamsRequest) LogStream(com.amazonaws.services.logs.model.LogStream) ListTablesResponse(com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse) DescribeLogStreamsResult(com.amazonaws.services.logs.model.DescribeLogStreamsResult) DescribeLogGroupsRequest(com.amazonaws.services.logs.model.DescribeLogGroupsRequest) Test(org.junit.Test)

Aggregations

DescribeLogStreamsRequest (com.amazonaws.services.logs.model.DescribeLogStreamsRequest)11 LogStream (com.amazonaws.services.logs.model.LogStream)11 DescribeLogStreamsResult (com.amazonaws.services.logs.model.DescribeLogStreamsResult)10 Test (org.junit.Test)4 InvocationOnMock (org.mockito.invocation.InvocationOnMock)4 TableName (com.amazonaws.athena.connector.lambda.domain.TableName)3 DescribeLogGroupsRequest (com.amazonaws.services.logs.model.DescribeLogGroupsRequest)3 ArrayList (java.util.ArrayList)3 CreateLogStreamRequest (com.amazonaws.services.logs.model.CreateLogStreamRequest)2 DescribeLogGroupsResult (com.amazonaws.services.logs.model.DescribeLogGroupsResult)2 LogGroup (com.amazonaws.services.logs.model.LogGroup)2 HashSet (java.util.HashSet)2 Test (org.junit.jupiter.api.Test)2 Level (ch.qos.logback.classic.Level)1 PatternLayout (ch.qos.logback.classic.PatternLayout)1 ILoggingEvent (ch.qos.logback.classic.spi.ILoggingEvent)1 LoggingEvent (ch.qos.logback.classic.spi.LoggingEvent)1 Block (com.amazonaws.athena.connector.lambda.data.Block)1 BlockAllocatorImpl (com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl)1 Constraints (com.amazonaws.athena.connector.lambda.domain.predicate.Constraints)1