MultipleCassandraClusterInstancesTestCase.java

/*
 * Copyright 2016 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 com.example;

import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.BeanClassLoaderAware;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.boot.autoconfigure.data.cassandra.CassandraDataAutoConfiguration;
import org.springframework.boot.test.SpringApplicationConfiguration;
import org.springframework.cassandra.config.CassandraCqlClusterFactoryBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.cassandra.config.CassandraEntityClassScanner;
import org.springframework.data.cassandra.config.CassandraSessionFactoryBean;
import org.springframework.data.cassandra.convert.CassandraConverter;
import org.springframework.data.cassandra.convert.MappingCassandraConverter;
import org.springframework.data.cassandra.mapping.BasicCassandraMappingContext;
import org.springframework.data.cassandra.mapping.CassandraMappingContext;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;

import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.Session;

/**
 * @author Mark Paluch
 */
@RunWith(SpringJUnit4ClassRunner.class)
@SpringApplicationConfiguration(classes = MultipleCassandraClusterInstancesTestCase.Config.class)
public class MultipleCassandraClusterInstancesTestCase {

    @Autowired
    Session session1;

    @Autowired
    Session session2;

    @Test
    public void test() {

        assert session1 != session2;
        assert session1.getCluster() != session2.getCluster();
    }

    @SpringBootApplication(exclude = { CassandraDataAutoConfiguration.class })
    @Configuration
    static class Config implements BeanClassLoaderAware {

        private ClassLoader classLoader;

        @Bean
        public CassandraCqlClusterFactoryBean cluster1() {

            CassandraCqlClusterFactoryBean bean = new CassandraCqlClusterFactoryBean();

            bean.setContactPoints("localhost");
            bean.setPort(9042);

            return bean;
        }

        @Bean
        public CassandraCqlClusterFactoryBean cluster2() {

            CassandraCqlClusterFactoryBean bean = new CassandraCqlClusterFactoryBean();

            bean.setContactPoints("127.0.0.1");
            bean.setPort(9042);

            return bean;
        }

        @Bean
        public CassandraSessionFactoryBean session1(@Qualifier("cluster1") Cluster cluster1,
                CassandraConverter cassandraConverter) {

            CassandraSessionFactoryBean bean = new CassandraSessionFactoryBean();

            bean.setCluster(cluster1);
            bean.setConverter(cassandraConverter);

            return bean;
        }

        @Bean
        public CassandraSessionFactoryBean session2(@Qualifier("cluster2") Cluster cluster2,
                CassandraConverter cassandraConverter) {

            CassandraSessionFactoryBean bean = new CassandraSessionFactoryBean();

            bean.setCluster(cluster2);
            bean.setConverter(cassandraConverter);

            return bean;
        }

        @Bean
        public CassandraMappingContext cassandraMapping() throws ClassNotFoundException {

            BasicCassandraMappingContext mappingContext = new BasicCassandraMappingContext();

            mappingContext.setBeanClassLoader(classLoader);
            mappingContext.setInitialEntitySet(CassandraEntityClassScanner.scan(getEntityBasePackages()));

            return mappingContext;
        }

        /**
         * Return the {@link CassandraConverter} instance to convert Rows to Objects, Objects to BuiltStatements
         */
        @Bean
        public CassandraConverter cassandraConverter() throws Exception {
            return new MappingCassandraConverter(cassandraMapping());
        }

        public Class<?> getEntityBasePackages() {
            return Config.class;
        }

        @Override
        public void setBeanClassLoader(ClassLoader classLoader) {
            this.classLoader = classLoader;
        }
    }
}