package com.datastax.driver.core;

import java.util.Iterator;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.testng.annotations.Test;

import static org.assertj.core.api.Assertions.assertThat;

import com.datastax.driver.core.policies.DelegatingLoadBalancingPolicy;
import com.datastax.driver.core.policies.LoadBalancingPolicy;
import com.datastax.driver.core.policies.Policies;
import com.datastax.driver.core.policies.ReconnectionPolicy;

public class ControlConnectionTest {
    @Test(groups = "short")
    public void should_prevent_simultaneous_reconnection_attempts() throws InterruptedException {
        CCMBridge ccm = null;
        Cluster cluster = null;

        // Custom load balancing policy that counts the number of calls to newQueryPlan().
        // Since we don't open any session from our Cluster, only the control connection reattempts are calling this
        // method, therefore the invocation count is equal to the number of attempts.
        QueryPlanCountingPolicy loadBalancingPolicy = new QueryPlanCountingPolicy(Policies.defaultLoadBalancingPolicy());
        AtomicInteger reconnectionAttempts = loadBalancingPolicy.counter;

        // Custom reconnection policy with a very large delay (longer than the test duration), to make sure we count
        // only the first reconnection attempt of each reconnection handler.
        ReconnectionPolicy reconnectionPolicy = new ReconnectionPolicy() {
            @Override
            public ReconnectionSchedule newSchedule() {
                return new ReconnectionSchedule() {
                    @Override
                    public long nextDelayMs() {
                        return 60 * 1000;
                    }
                };
            }
        };

        try {
            ccm = CCMBridge.create("test", 2);
            // We pass only the first host as contact point, so we know the control connection will be on this host
            cluster = Cluster.builder()
                .addContactPoint(CCMBridge.ipOfNode(1))
                .withReconnectionPolicy(reconnectionPolicy)
                .withLoadBalancingPolicy(loadBalancingPolicy)
                .build();
            cluster.init();

            // Kill the control connection host, there should be exactly one reconnection attempt
            ccm.stop(1);
            TimeUnit.SECONDS.sleep(1); // Sleep for a while to make sure our final count is not the result of lucky timing
            assertThat(reconnectionAttempts.get()).isEqualTo(1);

            ccm.stop(2);
            TimeUnit.SECONDS.sleep(1);
            assertThat(reconnectionAttempts.get()).isEqualTo(2);

        } finally {
            if (cluster != null)
                cluster.close();
            if (ccm != null)
                ccm.remove();
        }
    }

    static class QueryPlanCountingPolicy extends DelegatingLoadBalancingPolicy {

        final AtomicInteger counter = new AtomicInteger();

        public QueryPlanCountingPolicy(LoadBalancingPolicy delegate) {
            super(delegate);
        }

        public Iterator<Host> newQueryPlan(String loggedKeyspace, Statement statement) {
            counter.incrementAndGet();
            return super.newQueryPlan(loggedKeyspace, statement);
        }
    }
}

@@ -2840,24 +2840,1627 @@
    %7D%0A    %7D%0A%0A
+    /**%0A     * Ensures that if the host that the Control Connection is connected to is removed/decommissioned that the%0A     * Control Connection is reestablished to another host.%0A     *%0A     * @since 2.0.9%0A     * @jira_ticket JAVA-597%0A     * @expected_result Control Connection is reestablished to another host.%0A     * @test_category control_connection%0A     */%0A    @Test(groups = %22long%22)%0A    public void should_reestablish_if_control_node_decommissioned() throws InterruptedException %7B%0A        CCMBridge ccm = null;%0A        Cluster cluster = null;%0A%0A        try %7B%0A            ccm = CCMBridge.create(%22test%22, 3);%0A%0A            cluster = Cluster.builder()%0A                    .addContactPoint(CCMBridge.ipOfNode(1))%0A                    .build();%0A            cluster.init();%0A%0A            // Ensure the control connection host is that of the first node.%0A            String controlHost = cluster.manager.controlConnection.connectedHost().getAddress().getHostAddress();%0A            assertThat(controlHost).isEqualTo(CCMBridge.ipOfNode(1));%0A%0A            // Decommission the node.%0A            ccm.decommissionNode(1);%0A%0A            // Ensure that the new control connection is not null and it's host is not equal to the decommissioned node.%0A            Host newHost = cluster.manager.controlConnection.connectedHost();%0A            assertThat(newHost).isNotNull();%0A            assertThat(newHost.getAddress().getHostAddress()).isNotEqualTo(controlHost);%0A        %7D finally %7B%0A            if (cluster != null)%0A                cluster.close();%0A            if (ccm != null)%0A                ccm.remove();%0A        %7D%0A    %7D%0A%0A
     static c

Add test for JAVA-597.