/*
 * Copyright DataStax, Inc.
 *
 * 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.datastax.oss.driver.internal.core.session;

import static com.datastax.oss.driver.Assertions.assertThat;
import static com.datastax.oss.driver.Assertions.assertThatStage;
import static org.awaitility.Awaitility.await;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.addresstranslation.AddressTranslator;
import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
import com.datastax.oss.driver.api.core.config.DriverConfig;
import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
import com.datastax.oss.driver.api.core.connection.ReconnectionPolicy;
import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance;
import com.datastax.oss.driver.api.core.metadata.Metadata;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.NodeState;
import com.datastax.oss.driver.api.core.metadata.NodeStateListener;
import com.datastax.oss.driver.api.core.metadata.schema.SchemaChangeListener;
import com.datastax.oss.driver.api.core.retry.RetryPolicy;
import com.datastax.oss.driver.api.core.session.Session;
import com.datastax.oss.driver.api.core.specex.SpeculativeExecutionPolicy;
import com.datastax.oss.driver.api.core.tracker.RequestTracker;
import com.datastax.oss.driver.internal.core.context.EventBus;
import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
import com.datastax.oss.driver.internal.core.context.NettyOptions;
import com.datastax.oss.driver.internal.core.control.ControlConnection;
import com.datastax.oss.driver.internal.core.metadata.DefaultEndPoint;
import com.datastax.oss.driver.internal.core.metadata.DefaultNode;
import com.datastax.oss.driver.internal.core.metadata.DistanceEvent;
import com.datastax.oss.driver.internal.core.metadata.LoadBalancingPolicyWrapper;
import com.datastax.oss.driver.internal.core.metadata.MetadataManager;
import com.datastax.oss.driver.internal.core.metadata.NodeStateEvent;
import com.datastax.oss.driver.internal.core.metadata.TestNodeFactory;
import com.datastax.oss.driver.internal.core.metadata.TopologyMonitor;
import com.datastax.oss.driver.internal.core.metrics.MetricsFactory;
import com.datastax.oss.driver.internal.core.pool.ChannelPool;
import com.datastax.oss.driver.internal.core.pool.ChannelPoolFactory;
import com.datastax.oss.driver.internal.core.util.concurrent.CompletableFutures;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
import io.netty.channel.DefaultEventLoopGroup;
import io.netty.util.concurrent.DefaultPromise;
import io.netty.util.concurrent.GlobalEventExecutor;
import java.time.Duration;
import java.util.Collections;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.TimeUnit;
import org.junit.Before;
import org.junit.Test;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
import org.mockito.verification.VerificationWithTimeout;

public class DefaultSessionPoolsTest {

  private static final CqlIdentifier KEYSPACE = CqlIdentifier.fromInternal("ks");
  /** How long we wait when verifying mocks for async invocations */
  protected static final VerificationWithTimeout VERIFY_TIMEOUT = timeout(500);

  @Mock private InternalDriverContext context;
  @Mock private NettyOptions nettyOptions;
  @Mock private ChannelPoolFactory channelPoolFactory;
  @Mock private MetadataManager metadataManager;
  @Mock private TopologyMonitor topologyMonitor;
  @Mock private LoadBalancingPolicyWrapper loadBalancingPolicyWrapper;
  @Mock private DriverConfigLoader configLoader;
  @Mock private Metadata metadata;
  @Mock private DriverConfig config;
  @Mock private DriverExecutionProfile defaultProfile;
  @Mock private ReconnectionPolicy reconnectionPolicy;
  @Mock private RetryPolicy retryPolicy;
  @Mock private SpeculativeExecutionPolicy speculativeExecutionPolicy;
  @Mock private AddressTranslator addressTranslator;
  @Mock private ControlConnection controlConnection;
  @Mock private MetricsFactory metricsFactory;
  @Mock private NodeStateListener nodeStateListener;
  @Mock private SchemaChangeListener schemaChangeListener;
  @Mock private RequestTracker requestTracker;

  private DefaultNode node1;
  private DefaultNode node2;
  private DefaultNode node3;
  private EventBus eventBus;

  @Before
  public void setup() {
    MockitoAnnotations.initMocks(this);

    DefaultEventLoopGroup adminEventLoopGroup = new DefaultEventLoopGroup(1);
    when(nettyOptions.adminEventExecutorGroup()).thenReturn(adminEventLoopGroup);
    when(context.getNettyOptions()).thenReturn(nettyOptions);

    // Config:
    when(defaultProfile.getBoolean(DefaultDriverOption.REQUEST_WARN_IF_SET_KEYSPACE))
        .thenReturn(true);
    when(defaultProfile.getBoolean(DefaultDriverOption.REPREPARE_ENABLED)).thenReturn(false);
    when(defaultProfile.isDefined(DefaultDriverOption.PROTOCOL_VERSION)).thenReturn(true);
    when(defaultProfile.getDuration(DefaultDriverOption.METADATA_TOPOLOGY_WINDOW))
        .thenReturn(Duration.ZERO);
    when(defaultProfile.getInt(DefaultDriverOption.METADATA_TOPOLOGY_MAX_EVENTS)).thenReturn(1);
    when(config.getDefaultProfile()).thenReturn(defaultProfile);
    when(context.getConfig()).thenReturn(config);

    // Init sequence:
    when(metadataManager.refreshNodes()).thenReturn(CompletableFuture.completedFuture(null));
    when(metadataManager.refreshSchema(null, false, true))
        .thenReturn(CompletableFuture.completedFuture(null));
    when(context.getMetadataManager()).thenReturn(metadataManager);

    when(topologyMonitor.init()).thenReturn(CompletableFuture.completedFuture(null));
    when(context.getTopologyMonitor()).thenReturn(topologyMonitor);

    when(context.getLoadBalancingPolicyWrapper()).thenReturn(loadBalancingPolicyWrapper);

    when(context.getConfigLoader()).thenReturn(configLoader);

    when(context.getMetricsFactory()).thenReturn(metricsFactory);

    // Runtime behavior:
    when(context.getSessionName()).thenReturn("test");

    when(context.getChannelPoolFactory()).thenReturn(channelPoolFactory);

    eventBus = spy(new EventBus("test"));
    when(context.getEventBus()).thenReturn(eventBus);

    node1 = mockLocalNode(1);
    node2 = mockLocalNode(2);
    node3 = mockLocalNode(3);
    @SuppressWarnings("ConstantConditions")
    ImmutableMap<UUID, Node> nodes =
        ImmutableMap.of(
            node1.getHostId(), node1,
            node2.getHostId(), node2,
            node3.getHostId(), node3);
    when(metadata.getNodes()).thenReturn(nodes);
    when(metadataManager.getMetadata()).thenReturn(metadata);

    PoolManager poolManager = new PoolManager(context);
    when(context.getPoolManager()).thenReturn(poolManager);

    // Shutdown sequence:
    when(context.getReconnectionPolicy()).thenReturn(reconnectionPolicy);
    when(context.getRetryPolicy(DriverExecutionProfile.DEFAULT_NAME)).thenReturn(retryPolicy);
    when(context.getSpeculativeExecutionPolicies())
        .thenReturn(
            ImmutableMap.of(DriverExecutionProfile.DEFAULT_NAME, speculativeExecutionPolicy));
    when(context.getAddressTranslator()).thenReturn(addressTranslator);
    when(context.getNodeStateListener()).thenReturn(nodeStateListener);
    when(context.getSchemaChangeListener()).thenReturn(schemaChangeListener);
    when(context.getRequestTracker()).thenReturn(requestTracker);

    when(metadataManager.closeAsync()).thenReturn(CompletableFuture.completedFuture(null));
    when(metadataManager.forceCloseAsync()).thenReturn(CompletableFuture.completedFuture(null));

    when(topologyMonitor.closeAsync()).thenReturn(CompletableFuture.completedFuture(null));
    when(topologyMonitor.forceCloseAsync()).thenReturn(CompletableFuture.completedFuture(null));

    when(context.getControlConnection()).thenReturn(controlConnection);
    when(controlConnection.closeAsync()).thenReturn(CompletableFuture.completedFuture(null));
    when(controlConnection.forceCloseAsync()).thenReturn(CompletableFuture.completedFuture(null));

    DefaultPromise<Void> nettyCloseFuture = new DefaultPromise<>(GlobalEventExecutor.INSTANCE);
    nettyCloseFuture.setSuccess(null);
    when(nettyOptions.onClose()).thenAnswer(invocation -> nettyCloseFuture);
  }

  @Test
  public void should_initialize_pools_with_distances() {
    when(node3.getDistance()).thenReturn(NodeDistance.REMOTE);

    CompletableFuture<ChannelPool> pool1Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool3Future = new CompletableFuture<>();
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .pending(node1, KEYSPACE, NodeDistance.LOCAL, pool1Future)
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .pending(node3, KEYSPACE, NodeDistance.REMOTE, pool3Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.REMOTE);

    assertThatStage(initFuture).isNotDone();

    pool1Future.complete(pool1);
    pool2Future.complete(pool2);
    pool3Future.complete(pool3);

    assertThatStage(initFuture)
        .isSuccess(
            session ->
                assertThat(((DefaultSession) session).getPools())
                    .containsValues(pool1, pool2, pool3));
  }

  @Test
  public void should_not_connect_to_ignored_nodes() {
    when(node2.getDistance()).thenReturn(NodeDistance.IGNORED);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // Initial connection
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture)
        .isSuccess(
            session ->
                assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_not_connect_to_forced_down_nodes() {
    when(node2.getState()).thenReturn(NodeState.FORCED_DOWN);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // Initial connection
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture)
        .isSuccess(
            session ->
                assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_adjust_distance_if_changed_while_init() {
    CompletableFuture<ChannelPool> pool1Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool3Future = new CompletableFuture<>();
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .pending(node1, KEYSPACE, NodeDistance.LOCAL, pool1Future)
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .pending(node3, KEYSPACE, NodeDistance.LOCAL, pool3Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);

    assertThatStage(initFuture).isNotDone();

    // Distance changes while init still pending
    eventBus.fire(new DistanceEvent(NodeDistance.REMOTE, node2));

    pool1Future.complete(pool1);
    pool2Future.complete(pool2);
    pool3Future.complete(pool3);

    verify(pool2, VERIFY_TIMEOUT).resize(NodeDistance.REMOTE);

    assertThatStage(initFuture)
        .isSuccess(
            session ->
                assertThat(((DefaultSession) session).getPools())
                    .containsValues(pool1, pool2, pool3));
  }

  @Test
  public void should_remove_pool_if_ignored_while_init() {
    CompletableFuture<ChannelPool> pool1Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool3Future = new CompletableFuture<>();
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .pending(node1, KEYSPACE, NodeDistance.LOCAL, pool1Future)
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .pending(node3, KEYSPACE, NodeDistance.LOCAL, pool3Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);

    assertThatStage(initFuture).isNotDone();

    // Distance changes while init still pending
    eventBus.fire(new DistanceEvent(NodeDistance.IGNORED, node2));

    pool1Future.complete(pool1);
    pool2Future.complete(pool2);
    pool3Future.complete(pool3);

    verify(pool2, VERIFY_TIMEOUT).closeAsync();

    assertThatStage(initFuture)
        .isSuccess(
            session ->
                assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_remove_pool_if_forced_down_while_init() {
    CompletableFuture<ChannelPool> pool1Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    CompletableFuture<ChannelPool> pool3Future = new CompletableFuture<>();
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .pending(node1, KEYSPACE, NodeDistance.LOCAL, pool1Future)
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .pending(node3, KEYSPACE, NodeDistance.LOCAL, pool3Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);

    assertThatStage(initFuture).isNotDone();

    // Forced down while init still pending
    eventBus.fire(NodeStateEvent.changed(NodeState.UP, NodeState.FORCED_DOWN, node2));

    pool1Future.complete(pool1);
    pool2Future.complete(pool2);
    pool3Future.complete(pool3);

    verify(pool2, VERIFY_TIMEOUT).closeAsync();

    assertThatStage(initFuture)
        .isSuccess(
            session ->
                assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_resize_pool_if_distance_changes() {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();

    eventBus.fire(new DistanceEvent(NodeDistance.REMOTE, node2));
    verify(pool2, timeout(500)).resize(NodeDistance.REMOTE);
  }

  @Test
  public void should_remove_pool_if_node_becomes_ignored() {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();

    eventBus.fire(new DistanceEvent(NodeDistance.IGNORED, node2));
    verify(pool2, timeout(500)).closeAsync();

    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    await()
        .untilAsserted(
            () -> assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_do_nothing_if_node_becomes_ignored_but_was_already_ignored()
      throws InterruptedException {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();

    eventBus.fire(new DistanceEvent(NodeDistance.IGNORED, node2));
    verify(pool2, timeout(100)).closeAsync();

    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    // Fire the same event again, nothing should happen
    eventBus.fire(new DistanceEvent(NodeDistance.IGNORED, node2));
    TimeUnit.MILLISECONDS.sleep(200);
    factoryHelper.verifyNoMoreCalls();
  }

  @Test
  public void should_recreate_pool_if_node_becomes_not_ignored() {
    when(node2.getDistance()).thenReturn(NodeDistance.IGNORED);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // Initial connection
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // When node2 becomes not ignored
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    eventBus.fire(new DistanceEvent(NodeDistance.LOCAL, node2));

    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    await()
        .untilAsserted(
            () ->
                assertThat(((DefaultSession) session).getPools())
                    .containsValues(pool1, pool2, pool3));
  }

  @Test
  public void should_remove_pool_if_node_is_forced_down() {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();

    eventBus.fire(NodeStateEvent.changed(NodeState.UP, NodeState.FORCED_DOWN, node2));
    verify(pool2, timeout(500)).closeAsync();

    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    await()
        .untilAsserted(
            () -> assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_recreate_pool_if_node_is_forced_back_up() {
    when(node2.getState()).thenReturn(NodeState.FORCED_DOWN);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // init
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // when node2 comes back up
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    eventBus.fire(NodeStateEvent.changed(NodeState.FORCED_DOWN, NodeState.UP, node2));
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    await()
        .untilAsserted(
            () ->
                assertThat(((DefaultSession) session).getPools())
                    .containsValues(pool1, pool2, pool3));
  }

  @Test
  public void should_not_recreate_pool_if_node_is_forced_back_up_but_ignored() {
    when(node2.getState()).thenReturn(NodeState.FORCED_DOWN);
    when(node2.getDistance()).thenReturn(NodeDistance.IGNORED);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // init
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    eventBus.fire(NodeStateEvent.changed(NodeState.FORCED_DOWN, NodeState.UP, node2));
    await()
        .untilAsserted(
            () -> assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
    factoryHelper.verifyNoMoreCalls();
  }

  @Test
  public void should_adjust_distance_if_changed_while_recreating() {
    when(node2.getDistance()).thenReturn(NodeDistance.IGNORED);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // Initial connection
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // When node2 becomes not ignored
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    eventBus.fire(new DistanceEvent(NodeDistance.LOCAL, node2));

    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);

    // Distance changes again while pool init is in progress
    eventBus.fire(new DistanceEvent(NodeDistance.REMOTE, node2));

    // Now pool init succeeds
    pool2Future.complete(pool2);

    // Pool should have been adjusted
    verify(pool2, VERIFY_TIMEOUT).resize(NodeDistance.REMOTE);
    await()
        .untilAsserted(
            () ->
                assertThat(((DefaultSession) session).getPools())
                    .containsValues(pool1, pool2, pool3));
  }

  @Test
  public void should_remove_pool_if_ignored_while_recreating() {
    when(node2.getDistance()).thenReturn(NodeDistance.IGNORED);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // Initial connection
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // When node2 becomes not ignored
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    eventBus.fire(new DistanceEvent(NodeDistance.LOCAL, node2));

    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);

    // Distance changes to ignored while pool init is in progress
    eventBus.fire(new DistanceEvent(NodeDistance.IGNORED, node2));

    // Now pool init succeeds
    pool2Future.complete(pool2);

    // Pool should have been closed
    verify(pool2, VERIFY_TIMEOUT).closeAsync();

    await()
        .untilAsserted(
            () -> assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_remove_pool_if_forced_down_while_recreating() {
    when(node2.getDistance()).thenReturn(NodeDistance.IGNORED);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // Initial connection
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // When node2 becomes not ignored
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    eventBus.fire(new DistanceEvent(NodeDistance.LOCAL, node2));

    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);

    // Forced down while pool init is in progress
    eventBus.fire(NodeStateEvent.changed(NodeState.UP, NodeState.FORCED_DOWN, node2));

    // Now pool init succeeds
    pool2Future.complete(pool2);

    // Pool should have been closed
    verify(pool2, VERIFY_TIMEOUT).closeAsync();
    await()
        .untilAsserted(
            () -> assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3));
  }

  @Test
  public void should_close_all_pools_when_closing() {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());

    CompletionStage<Void> closeFuture = session.closeAsync();
    assertThatStage(closeFuture).isSuccess();

    verify(pool1, VERIFY_TIMEOUT).closeAsync();
    verify(pool2, VERIFY_TIMEOUT).closeAsync();
    verify(pool3, VERIFY_TIMEOUT).closeAsync();
  }

  @Test
  public void should_force_close_all_pools_when_force_closing() {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());

    CompletionStage<Void> closeFuture = session.forceCloseAsync();
    assertThatStage(closeFuture).isSuccess();

    verify(pool1, VERIFY_TIMEOUT).forceCloseAsync();
    verify(pool2, VERIFY_TIMEOUT).forceCloseAsync();
    verify(pool3, VERIFY_TIMEOUT).forceCloseAsync();
  }

  @Test
  public void should_close_pool_if_recreated_while_closing() {
    when(node2.getState()).thenReturn(NodeState.FORCED_DOWN);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // init
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // when node2 comes back up
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(((DefaultSession) session).getPools()).containsValues(pool1, pool3);

    // node2 comes back up, start initializing a pool for it
    eventBus.fire(NodeStateEvent.changed(NodeState.FORCED_DOWN, NodeState.UP, node2));
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);

    // but the session gets closed before pool init completes
    CompletionStage<Void> closeFuture = session.closeAsync();
    assertThatStage(closeFuture).isSuccess();

    // now pool init completes
    pool2Future.complete(pool2);

    // Pool should have been closed
    verify(pool2, VERIFY_TIMEOUT).forceCloseAsync();
  }

  @Test
  public void should_set_keyspace_on_all_pools() {
    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node2, KEYSPACE, NodeDistance.LOCAL, pool2)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    Session session = CompletableFutures.getCompleted(initFuture.toCompletableFuture());

    CqlIdentifier newKeyspace = CqlIdentifier.fromInternal("newKeyspace");
    ((DefaultSession) session).setKeyspace(newKeyspace);

    verify(pool1, VERIFY_TIMEOUT).setKeyspace(newKeyspace);
    verify(pool2, VERIFY_TIMEOUT).setKeyspace(newKeyspace);
    verify(pool3, VERIFY_TIMEOUT).setKeyspace(newKeyspace);
  }

  @Test
  public void should_set_keyspace_on_pool_if_recreated_while_switching_keyspace() {
    when(node2.getState()).thenReturn(NodeState.FORCED_DOWN);

    ChannelPool pool1 = mockPool(node1);
    ChannelPool pool2 = mockPool(node2);
    CompletableFuture<ChannelPool> pool2Future = new CompletableFuture<>();
    ChannelPool pool3 = mockPool(node3);
    MockChannelPoolFactoryHelper factoryHelper =
        MockChannelPoolFactoryHelper.builder(channelPoolFactory)
            // init
            .success(node1, KEYSPACE, NodeDistance.LOCAL, pool1)
            .success(node3, KEYSPACE, NodeDistance.LOCAL, pool3)
            // when node2 comes back up
            .pending(node2, KEYSPACE, NodeDistance.LOCAL, pool2Future)
            .build();

    CompletionStage<CqlSession> initFuture = newSession();

    factoryHelper.waitForCall(node1, KEYSPACE, NodeDistance.LOCAL);
    factoryHelper.waitForCall(node3, KEYSPACE, NodeDistance.LOCAL);
    assertThatStage(initFuture).isSuccess();
    DefaultSession session =
        (DefaultSession) CompletableFutures.getCompleted(initFuture.toCompletableFuture());
    assertThat(session.getPools()).containsValues(pool1, pool3);

    // node2 comes back up, start initializing a pool for it
    eventBus.fire(NodeStateEvent.changed(NodeState.FORCED_DOWN, NodeState.UP, node2));
    factoryHelper.waitForCall(node2, KEYSPACE, NodeDistance.LOCAL);

    // Keyspace gets changed on the session in the meantime, node2's pool will miss it
    CqlIdentifier newKeyspace = CqlIdentifier.fromInternal("newKeyspace");
    session.setKeyspace(newKeyspace);
    verify(pool1, VERIFY_TIMEOUT).setKeyspace(newKeyspace);
    verify(pool3, VERIFY_TIMEOUT).setKeyspace(newKeyspace);

    // now pool init completes
    pool2Future.complete(pool2);

    // Pool should have been closed
    verify(pool2, VERIFY_TIMEOUT).setKeyspace(newKeyspace);
  }

  private ChannelPool mockPool(Node node) {
    ChannelPool pool = mock(ChannelPool.class);
    when(pool.getNode()).thenReturn(node);
    when(pool.getInitialKeyspaceName()).thenReturn(KEYSPACE);
    when(pool.setKeyspace(any(CqlIdentifier.class)))
        .thenReturn(CompletableFuture.completedFuture(null));
    CompletableFuture<Void> closeFuture = new CompletableFuture<>();
    when(pool.closeFuture()).thenReturn(closeFuture);
    when(pool.closeAsync())
        .then(
            i -> {
              closeFuture.complete(null);
              return closeFuture;
            });
    when(pool.forceCloseAsync())
        .then(
            i -> {
              closeFuture.complete(null);
              return closeFuture;
            });
    return pool;
  }

  private CompletionStage<CqlSession> newSession() {
    return DefaultSession.init(context, Collections.emptySet(), KEYSPACE);
  }

  private static DefaultNode mockLocalNode(int i) {
    DefaultNode node = mock(DefaultNode.class);
    when(node.getHostId()).thenReturn(UUID.randomUUID());
    DefaultEndPoint endPoint = TestNodeFactory.newEndPoint(i);
    when(node.getEndPoint()).thenReturn(endPoint);
    when(node.getBroadcastRpcAddress()).thenReturn(Optional.of(endPoint.resolve()));
    when(node.getDistance()).thenReturn(NodeDistance.LOCAL);
    when(node.toString()).thenReturn("node" + i);
    return node;
  }
}
