Search in sources :

Example 11 with GetSecretValueResult

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

the class OracleMetadataHandlerTest method setup.

@Before
public void setup() {
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class, Mockito.RETURNS_DEEP_STUBS);
    this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
    this.oracleMetadataHandler = new OracleMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory);
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Connection(java.sql.Connection) OracleMetadataHandler(com.amazonaws.athena.connectors.oracle.OracleMetadataHandler) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Example 12 with GetSecretValueResult

use of com.amazonaws.services.secretsmanager.model.GetSecretValueResult 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());
    });
}
Also used : GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) InvocationOnMock(org.mockito.invocation.InvocationOnMock) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) LocalKeyFactory(com.amazonaws.athena.connector.lambda.security.LocalKeyFactory) Before(org.junit.Before)

Example 13 with GetSecretValueResult

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

the class RedisRecordHandlerTest method setUp.

@Before
public void setUp() {
    logger.info("{}: enter", testName.getMethodName());
    when(mockFactory.getOrCreateConn(eq(decodedEndpoint), anyBoolean(), anyBoolean(), anyString())).thenReturn(mockConnection);
    when(mockConnection.sync()).thenReturn(mockSyncCommands);
    allocator = new BlockAllocatorImpl();
    amazonS3 = mock(AmazonS3.class);
    when(amazonS3.putObject(anyObject(), anyObject(), anyObject(), anyObject())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        InputStream inputStream = (InputStream) invocationOnMock.getArguments()[2];
        ByteHolder byteHolder = new ByteHolder();
        byteHolder.setBytes(ByteStreams.toByteArray(inputStream));
        synchronized (mockS3Storage) {
            mockS3Storage.add(byteHolder);
            logger.info("puObject: total size " + mockS3Storage.size());
        }
        return mock(PutObjectResult.class);
    });
    when(amazonS3.getObject(anyString(), anyString())).thenAnswer((InvocationOnMock invocationOnMock) -> {
        S3Object mockObject = mock(S3Object.class);
        ByteHolder byteHolder;
        synchronized (mockS3Storage) {
            byteHolder = mockS3Storage.get(0);
            mockS3Storage.remove(0);
            logger.info("getObject: total size " + mockS3Storage.size());
        }
        when(mockObject.getObjectContent()).thenReturn(new S3ObjectInputStream(new ByteArrayInputStream(byteHolder.getBytes()), null));
        return mockObject;
    });
    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());
    });
    handler = new RedisRecordHandler(amazonS3, mockSecretsManager, mockAthena, mockFactory);
    spillReader = new S3BlockSpillReader(amazonS3, allocator);
    logger.info("setUpBefore - exit");
}
Also used : AmazonS3(com.amazonaws.services.s3.AmazonS3) ByteArrayInputStream(java.io.ByteArrayInputStream) S3ObjectInputStream(com.amazonaws.services.s3.model.S3ObjectInputStream) InputStream(java.io.InputStream) S3ObjectInputStream(com.amazonaws.services.s3.model.S3ObjectInputStream) S3BlockSpillReader(com.amazonaws.athena.connector.lambda.data.S3BlockSpillReader) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) BlockAllocatorImpl(com.amazonaws.athena.connector.lambda.data.BlockAllocatorImpl) ByteArrayInputStream(java.io.ByteArrayInputStream) InvocationOnMock(org.mockito.invocation.InvocationOnMock) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) S3Object(com.amazonaws.services.s3.model.S3Object) Before(org.junit.Before)

Example 14 with GetSecretValueResult

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

the class SnowflakeMetadataHandlerTest method setup.

@Before
public void setup() {
    environmentVariables.set("pagecount", "300000");
    environmentVariables.set("partitionlimit", "300000");
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class, Mockito.RETURNS_DEEP_STUBS);
    this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"username\": \"testUser\", \"password\": \"testPassword\"}"));
    this.snowflakeMetadataHandler = new SnowflakeMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory);
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
    this.blockAllocator = Mockito.mock(BlockAllocator.class);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) SnowflakeMetadataHandler(com.amazonaws.athena.connectors.snowflake.SnowflakeMetadataHandler) Before(org.junit.Before)

Example 15 with GetSecretValueResult

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

the class SynapseMetadataHandlerTest method setup.

@Before
public void setup() {
    System.setProperty("aws.region", "us-east-1");
    this.jdbcConnectionFactory = Mockito.mock(JdbcConnectionFactory.class, Mockito.RETURNS_DEEP_STUBS);
    this.connection = Mockito.mock(Connection.class, Mockito.RETURNS_DEEP_STUBS);
    logger.info(" this.connection.." + this.connection);
    Mockito.when(this.jdbcConnectionFactory.getConnection(Mockito.any(JdbcCredentialProvider.class))).thenReturn(this.connection);
    this.secretsManager = Mockito.mock(AWSSecretsManager.class);
    this.athena = Mockito.mock(AmazonAthena.class);
    Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(new GetSecretValueRequest().withSecretId("testSecret")))).thenReturn(new GetSecretValueResult().withSecretString("{\"user\": \"testUser\", \"password\": \"testPassword\"}"));
    this.synapseMetadataHandler = new SynapseMetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory);
    this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
}
Also used : JdbcConnectionFactory(com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory) GetSecretValueResult(com.amazonaws.services.secretsmanager.model.GetSecretValueResult) FederatedIdentity(com.amazonaws.athena.connector.lambda.security.FederatedIdentity) AWSSecretsManager(com.amazonaws.services.secretsmanager.AWSSecretsManager) Connection(java.sql.Connection) GetSecretValueRequest(com.amazonaws.services.secretsmanager.model.GetSecretValueRequest) JdbcCredentialProvider(com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider) AmazonAthena(com.amazonaws.services.athena.AmazonAthena) Before(org.junit.Before)

Aggregations

GetSecretValueResult (com.amazonaws.services.secretsmanager.model.GetSecretValueResult)60 GetSecretValueRequest (com.amazonaws.services.secretsmanager.model.GetSecretValueRequest)51 AWSSecretsManager (com.amazonaws.services.secretsmanager.AWSSecretsManager)25 Before (org.junit.Before)21 JdbcConnectionFactory (com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory)18 JdbcCredentialProvider (com.amazonaws.athena.connectors.jdbc.connection.JdbcCredentialProvider)18 AmazonAthena (com.amazonaws.services.athena.AmazonAthena)17 FederatedIdentity (com.amazonaws.athena.connector.lambda.security.FederatedIdentity)16 Connection (java.sql.Connection)13 Test (org.junit.jupiter.api.Test)10 Test (org.junit.Test)8 AmazonS3 (com.amazonaws.services.s3.AmazonS3)6 DatabaseConnectionConfig (com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig)5 InvocationOnMock (org.mockito.invocation.InvocationOnMock)5 Map (java.util.Map)4 ParameterizedTest (org.junit.jupiter.params.ParameterizedTest)4 QueryStatusChecker (com.amazonaws.athena.connector.lambda.QueryStatusChecker)3 JsonProcessingException (com.fasterxml.jackson.core.JsonProcessingException)3 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)3 IOException (java.io.IOException)3