AmazonRdsReadReplicaAwareDataSourceFactoryBeanTest.java

/*
 * Copyright 2013-2014 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.springframework.cloud.aws.jdbc.rds;

import com.amazonaws.services.rds.AmazonRDS;
import com.amazonaws.services.rds.model.DBInstance;
import com.amazonaws.services.rds.model.DescribeDBInstancesRequest;
import com.amazonaws.services.rds.model.DescribeDBInstancesResult;
import com.amazonaws.services.rds.model.Endpoint;
import org.junit.Test;
import org.springframework.cloud.aws.jdbc.datasource.DataSourceFactory;
import org.springframework.cloud.aws.jdbc.datasource.DataSourceInformation;
import org.springframework.cloud.aws.jdbc.datasource.ReadOnlyRoutingDataSource;
import org.springframework.cloud.aws.jdbc.datasource.support.DatabaseType;
import org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy;

import javax.sql.DataSource;
import java.sql.Connection;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

/**
 * @author Agim Emruli
 */
public class AmazonRdsReadReplicaAwareDataSourceFactoryBeanTest {

    @Test
    public void afterPropertiesSet_instanceWithoutReadReplica_createsNoDataSourceRouter() throws Exception {
        //Arrange
        AmazonRDS amazonRDS = mock(AmazonRDS.class);
        DataSourceFactory dataSourceFactory = mock(DataSourceFactory.class);

        when(amazonRDS.describeDBInstances(new DescribeDBInstancesRequest().withDBInstanceIdentifier("test"))).thenReturn(
                new DescribeDBInstancesResult().
                        withDBInstances(new DBInstance().
                                withDBInstanceStatus("available").
                                withDBName("test").
                                withDBInstanceIdentifier("test").
                                withEngine("mysql").
                                withMasterUsername("admin").
                                withEndpoint(new Endpoint().
                                        withAddress("localhost").
                                        withPort(3306)
                                )
                        )
        );

        AmazonRdsReadReplicaAwareDataSourceFactoryBean amazonRdsDataSourceFactoryBean = new AmazonRdsReadReplicaAwareDataSourceFactoryBean(amazonRDS, "test", "secret");
        amazonRdsDataSourceFactoryBean.setDataSourceFactory(dataSourceFactory);
        when(dataSourceFactory.createDataSource(new DataSourceInformation(DatabaseType.MYSQL, "localhost", 3306, "test", "admin", "secret"))).thenReturn(mock(DataSource.class));

        //Act
        amazonRdsDataSourceFactoryBean.afterPropertiesSet();

        //Assert
        DataSource datasource = amazonRdsDataSourceFactoryBean.getObject();
        assertNotNull(datasource);

        verify(dataSourceFactory, times(1)).createDataSource(new DataSourceInformation(DatabaseType.MYSQL, "localhost", 3306, "test", "admin", "secret"));
    }

    @Test
    public void afterPropertiesSet_instanceWithReadReplica_createsDataSourceRouter() throws Exception {
        //Arrange
        AmazonRDS amazonRDS = mock(AmazonRDS.class);
        DataSourceFactory dataSourceFactory = mock(DataSourceFactory.class);


        when(amazonRDS.describeDBInstances(new DescribeDBInstancesRequest().withDBInstanceIdentifier("test"))).thenReturn(
                new DescribeDBInstancesResult().
                        withDBInstances(new DBInstance().
                                withDBInstanceStatus("available").
                                withDBName("test").
                                withDBInstanceIdentifier("test").
                                withEngine("mysql").
                                withMasterUsername("admin").
                                withEndpoint(new Endpoint().
                                        withAddress("localhost").
                                        withPort(3306)
                                ).withReadReplicaDBInstanceIdentifiers("read1", "read2")
                        )
        );

        when(amazonRDS.describeDBInstances(new DescribeDBInstancesRequest().withDBInstanceIdentifier("read1"))).thenReturn(
                new DescribeDBInstancesResult().
                        withDBInstances(new DBInstance().
                                withDBInstanceStatus("available").
                                withDBName("read1").
                                withDBInstanceIdentifier("read1").
                                withEngine("mysql").
                                withMasterUsername("admin").
                                withEndpoint(new Endpoint().
                                        withAddress("localhost").
                                        withPort(3306)
                                )
                        )
        );

        when(amazonRDS.describeDBInstances(new DescribeDBInstancesRequest().withDBInstanceIdentifier("read2"))).thenReturn(
                new DescribeDBInstancesResult().
                        withDBInstances(new DBInstance().
                                withDBInstanceStatus("available").
                                withDBName("read2").
                                withDBInstanceIdentifier("read2").
                                withEngine("mysql").
                                withMasterUsername("admin").
                                withEndpoint(new Endpoint().
                                        withAddress("localhost").
                                        withPort(3306)
                                )
                        )
        );

        DataSource createdDataSource = mock(DataSource.class);
        Connection connection = mock(Connection.class);

        when(dataSourceFactory.createDataSource(new DataSourceInformation(DatabaseType.MYSQL, "localhost", 3306, "test", "admin", "secret"))).thenReturn(createdDataSource);
        when(dataSourceFactory.createDataSource(new DataSourceInformation(DatabaseType.MYSQL, "localhost", 3306, "read1", "admin", "secret"))).thenReturn(createdDataSource);
        when(dataSourceFactory.createDataSource(new DataSourceInformation(DatabaseType.MYSQL, "localhost", 3306, "read2", "admin", "secret"))).thenReturn(createdDataSource);
        when(createdDataSource.getConnection()).thenReturn(connection);

        AmazonRdsReadReplicaAwareDataSourceFactoryBean amazonRdsDataSourceFactoryBean = new AmazonRdsReadReplicaAwareDataSourceFactoryBean(amazonRDS, "test", "secret");
        amazonRdsDataSourceFactoryBean.setDataSourceFactory(dataSourceFactory);

        //Act
        amazonRdsDataSourceFactoryBean.afterPropertiesSet();

        //Assert
        DataSource datasource = amazonRdsDataSourceFactoryBean.getObject();
        assertNotNull(datasource);
        assertTrue(datasource instanceof LazyConnectionDataSourceProxy);

        ReadOnlyRoutingDataSource source = (ReadOnlyRoutingDataSource) ((LazyConnectionDataSourceProxy) datasource).getTargetDataSource();
        assertEquals(3, source.getDataSources().size());
    }
}