/**
 * Copyright 2016 Palantir Technologies
 *
 * Licensed under the BSD-3 License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://opensource.org/licenses/BSD-3-Clause
 *
 * 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.palantir.atlasdb.keyvalue.cassandra;

import static org.hamcrest.core.IsEqual.equalTo;
import static org.junit.Assert.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

import java.net.InetSocketAddress;
import java.net.UnknownHostException;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.junit.Test;

import com.google.common.collect.ImmutableSet;
import com.palantir.atlasdb.cassandra.CassandraKeyValueServiceConfig;

public class CassandraClientPoolUnitTest {
    public static final int POOL_REFRESH_INTERVAL_SECONDS = 10;
    public static final int DEFAULT_PORT = 5000;
    public static final int OTHER_PORT = 6000;
    public static final String HOST_1 = "1.0.0.0";
    public static final String HOST_2 = "2.0.0.0";
    public static final String HOST_3 = "3.0.0.0";

    @Test
    public void shouldReturnAddressForSingleHostInPool() throws UnknownHostException {
        InetSocketAddress host = new InetSocketAddress(HOST_1, DEFAULT_PORT);
        CassandraClientPool cassandraClientPool = clientPoolWithServersInCurrentPool(ImmutableSet.of(host));

        InetSocketAddress resolvedHost = cassandraClientPool.getAddressForHost(host.getHostName());

        assertThat(resolvedHost, equalTo(host));
    }

    @Test
    public void shouldReturnAddressForSingleServer() throws UnknownHostException {
        InetSocketAddress host = new InetSocketAddress(HOST_1, DEFAULT_PORT);
        CassandraClientPool cassandraClientPool = clientPoolWithServers(ImmutableSet.of(host));

        InetSocketAddress resolvedHost = cassandraClientPool.getAddressForHost(host.getHostName());

        assertThat(resolvedHost, equalTo(host));
    }

    @Test
    public void shouldUseCommonPortIfThereIsOnlyOneAndNoAddressMatches() throws UnknownHostException {
        InetSocketAddress host1 = new InetSocketAddress(HOST_1, DEFAULT_PORT);
        InetSocketAddress host2 = new InetSocketAddress(HOST_2, DEFAULT_PORT);
        CassandraClientPool cassandraClientPool = clientPoolWithServers(ImmutableSet.of(host1, host2));

        InetSocketAddress resolvedHost = cassandraClientPool.getAddressForHost(HOST_3);

        assertThat(resolvedHost, equalTo(new InetSocketAddress(HOST_3, DEFAULT_PORT)));
    }


    @Test(expected = UnknownHostException.class)
    public void shouldThrowIfPortsAreNotTheSameAddressDoesNotMatch() throws UnknownHostException {
        InetSocketAddress host1 = new InetSocketAddress(HOST_1, DEFAULT_PORT);
        InetSocketAddress host2 = new InetSocketAddress(HOST_2, OTHER_PORT);

        CassandraClientPool cassandraClientPool = clientPoolWithServers(ImmutableSet.of(host1, host2));

        InetSocketAddress resolvedHost = cassandraClientPool.getAddressForHost(HOST_3);

        assertThat(resolvedHost, equalTo(host1));
    }

    private CassandraClientPool clientPoolWithServers(ImmutableSet<InetSocketAddress> servers) {
        return clientPoolWith(servers, ImmutableSet.of());
    }

    private CassandraClientPool clientPoolWithServersInCurrentPool(ImmutableSet<InetSocketAddress> servers) {
        return clientPoolWith(ImmutableSet.of(), servers);
    }

    private CassandraClientPool clientPoolWith(
            ImmutableSet<InetSocketAddress> servers,
            ImmutableSet<InetSocketAddress> serversInPool) {
        CassandraKeyValueServiceConfig config = mock(CassandraKeyValueServiceConfig.class);
        when(config.poolRefreshIntervalSeconds()).thenReturn(POOL_REFRESH_INTERVAL_SECONDS);
        when(config.servers()).thenReturn(servers);

        CassandraClientPool cassandraClientPool = new CassandraClientPool(config);
        cassandraClientPool.currentPools = serversInPool.stream()
                .collect(Collectors.toMap(Function.identity(), address -> mock(CassandraClientPoolingContainer.class)));
        return cassandraClientPool;
    }
}

@@ -1331,16 +1331,20 @@
 ing HOST
+NAME
 _1 = %221.
@@ -1386,16 +1386,20 @@
 ing HOST
+NAME
 _2 = %222.
@@ -1441,16 +1441,20 @@
 ing HOST
+NAME
 _3 = %223.
@@ -1610,32 +1610,36 @@
 cketAddress(HOST
+NAME
 _1, DEFAULT_PORT
@@ -1830,34 +1830,26 @@
 ForHost(
-host.getHostName()
+HOSTNAME_1
 );%0A%0A    
@@ -2044,32 +2044,36 @@
 cketAddress(HOST
+NAME
 _1, DEFAULT_PORT
@@ -2255,26 +2255,18 @@
 ost(
-host.getHostName()
+HOSTNAME_1
 );%0A%0A
@@ -2486,32 +2486,36 @@
 cketAddress(HOST
+NAME
 _1, DEFAULT_PORT
@@ -2573,24 +2573,28 @@
 Address(HOST
+NAME
 _2, DEFAULT_
@@ -2780,32 +2780,36 @@
 ressForHost(HOST
+NAME
 _3);%0A%0A        as
@@ -2865,16 +2865,20 @@
 ess(HOST
+NAME
 _3, DEFA
@@ -3106,16 +3106,20 @@
 ess(HOST
+NAME
 _1, DEFA
@@ -3189,16 +3189,20 @@
 ess(HOST
+NAME
 _2, OTHE
@@ -3327,41 +3327,8 @@
     
- InetSocketAddress resolvedHost =
  cas
@@ -3370,61 +3370,14 @@
 HOST
-_3);%0A%0A        assertThat(resolvedHost, equalTo(host1)
+NAME_3
 );%0A 

Clean up ClientPoolUnitTest