use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.
the class ExampleMetadataHandlerTest method setUp.
@Before
public void setUp() {
logger.info("setUpBefore - enter");
allocator = new BlockAllocatorImpl();
metadataHandler = new ExampleMetadataHandler(new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "spill-bucket", "spill-prefix");
logger.info("setUpBefore - exit");
}
use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.
the class GlueMetadataHandlerTest method setUp.
@Before
public void setUp() throws Exception {
logger.info("{}: enter", testName.getMethodName());
handler = new GlueMetadataHandler(mockGlue, new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "glue-test", "spill-bucket", "spill-prefix") {
@Override
public GetTableLayoutResponse doGetTableLayout(BlockAllocator blockAllocator, GetTableLayoutRequest request) {
throw new UnsupportedOperationException();
}
@Override
public void getPartitions(BlockWriter blockWriter, GetTableLayoutRequest request, QueryStatusChecker queryStatusChecker) throws Exception {
throw new UnsupportedOperationException();
}
@Override
public GetSplitsResponse doGetSplits(BlockAllocator blockAllocator, GetSplitsRequest request) {
throw new UnsupportedOperationException();
}
};
allocator = new BlockAllocatorImpl();
// doListTables pagination.
when(mockGlue.getTables(any(GetTablesRequest.class))).thenAnswer((InvocationOnMock invocationOnMock) -> {
GetTablesRequest request = (GetTablesRequest) invocationOnMock.getArguments()[0];
String nextToken = request.getNextToken();
int pageSize = request.getMaxResults() == null ? UNLIMITED_PAGE_SIZE_VALUE : request.getMaxResults();
assertEquals(accountId, request.getCatalogId());
assertEquals(schema, request.getDatabaseName());
GetTablesResult mockResult = mock(GetTablesResult.class);
if (pageSize == UNLIMITED_PAGE_SIZE_VALUE) {
// Simulate full list of tables returned from Glue.
when(mockResult.getTableList()).thenReturn(unPaginatedTables);
when(mockResult.getNextToken()).thenReturn(null);
} else {
// Simulate paginated list of tables returned from Glue.
List<Table> paginatedTables = unPaginatedTables.stream().sorted(Comparator.comparing(Table::getName)).filter(table -> nextToken == null || table.getName().compareTo(nextToken) >= 0).limit(pageSize + 1).collect(Collectors.toList());
if (paginatedTables.size() > pageSize) {
when(mockResult.getNextToken()).thenReturn(paginatedTables.get(pageSize).getName());
when(mockResult.getTableList()).thenReturn(paginatedTables.subList(0, pageSize));
} else {
when(mockResult.getNextToken()).thenReturn(null);
when(mockResult.getTableList()).thenReturn(paginatedTables);
}
}
return mockResult;
});
}
use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.
the class NeptuneMetadataHandlerTest method setUp.
@Before
public void setUp() throws Exception {
logger.info("setUpBefore - enter");
allocator = new BlockAllocatorImpl();
handler = new NeptuneMetadataHandler(glue, neptuneConnection, new LocalKeyFactory(), mock(AWSSecretsManager.class), mock(AmazonAthena.class), "spill-bucket", "spill-prefix");
logger.info("setUpBefore - exit");
}
use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.
the class RedisMetadataHandlerTest method setUp.
@Before
public void setUp() throws Exception {
logger.info("{}: enter", testName.getMethodName());
when(mockFactory.getOrCreateConn(eq(decodedEndpoint), anyBoolean(), anyBoolean(), anyString())).thenReturn(mockConnection);
when(mockConnection.sync()).thenReturn(mockSyncCommands);
handler = new RedisMetadataHandler(mockGlue, new LocalKeyFactory(), mockSecretsManager, mockAthena, mockFactory, "bucket", "prefix");
allocator = new BlockAllocatorImpl();
when(mockSecretsManager.getSecretValue(any(GetSecretValueRequest.class))).thenAnswer((InvocationOnMock invocation) -> {
GetSecretValueRequest request = invocation.getArgumentAt(0, GetSecretValueRequest.class);
if ("endpoint".equalsIgnoreCase(request.getSecretId())) {
return new GetSecretValueResult().withSecretString(decodedEndpoint);
}
throw new RuntimeException("Unknown secret " + request.getSecretId());
});
}
use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.
the class TimestreamMetadataHandlerTest method setUp.
@Before
public void setUp() throws Exception {
handler = new TimestreamMetadataHandler(mockTsQuery, mockTsMeta, mockGlue, new LocalKeyFactory(), mockSecretsManager, mockAthena, "spillBucket", "spillPrefix");
allocator = new BlockAllocatorImpl();
}
Aggregations