/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.driver.core;

import static com.datastax.driver.core.Assertions.assertThat;
import static com.datastax.driver.core.HostDistance.LOCAL;
import static com.datastax.driver.core.TestUtils.nonDebouncingQueryOptions;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.scassandra.http.client.ClosedConnectionReport.CloseType.CLOSE;

import com.datastax.driver.core.policies.ConstantReconnectionPolicy;
import com.google.common.util.concurrent.Uninterruptibles;
import java.net.InetSocketAddress;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

public class HostConnectionPoolMultiTest {

  private ScassandraCluster scassandra;

  private Cluster cluster;

  @BeforeMethod(groups = {"short", "long"})
  private void setUp() {
    scassandra = ScassandraCluster.builder().withNodes(2).build();
    scassandra.init();
  }

  @AfterMethod(
      groups = {"short", "long"},
      alwaysRun = true)
  private void tearDown() {
    if (cluster != null) {
      cluster.close();
    }
    scassandra.stop();
  }

  private void createCluster(int core, int max) {
    PoolingOptions poolingOptions = new PoolingOptions().setConnectionsPerHost(LOCAL, core, max);
    SocketOptions socketOptions = new SocketOptions().setReadTimeoutMillis(1000);
    cluster =
        Cluster.builder()
            .addContactPoints(scassandra.address(1).getAddress())
            .withPort(scassandra.getBinaryPort())
            .withQueryOptions(nonDebouncingQueryOptions())
            .withPoolingOptions(poolingOptions)
            .withSocketOptions(socketOptions)
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    cluster.connect();
  }

  /**
   * Ensures that if all connections fail to a host on pool init that the host is marked down.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_mark_host_down_if_all_connections_fail_on_init() {
    // Prevent any connections on node 2.
    scassandra.node(2).currentClient().disableListener();
    createCluster(8, 8);

    // Node 2 should be in a down state while node 1 stays up.
    assertThat(cluster).host(2).goesDownWithin(10, SECONDS);
    assertThat(cluster).host(1).isUp();

    // Node 2 should come up as soon as it is able to reconnect.
    scassandra.node(2).currentClient().enableListener();
    assertThat(cluster).host(2).comesUpWithin(2, SECONDS);
  }

  /**
   * Ensures that if the control connection goes down, but the Host bound the control connection
   * still has an up pool, the Host should remain up and the Control Connection should be replaced.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_replace_control_connection_if_it_goes_down_but_host_remains_up() {
    createCluster(1, 2);

    // Ensure control connection is on node 1.
    assertThat(cluster).usesControlHost(1);

    // Identify the socket associated with the control connection.
    Connection controlConnection = cluster.manager.controlConnection.connectionRef.get();
    InetSocketAddress controlSocket = (InetSocketAddress) controlConnection.channel.localAddress();

    // Close the control connection.
    scassandra.node(1).currentClient().closeConnection(CLOSE, controlSocket);

    // Sleep reconnect interval * 2 to allow time to reconnect.
    Uninterruptibles.sleepUninterruptibly(2, SECONDS);

    // Ensure the control connection was replaced and host 1 remains up.
    assertThat(cluster).hasOpenControlConnection().host(1).isUp();
    assertThat(cluster.manager.controlConnection.connectionRef.get())
        .isNotEqualTo(controlConnection);
  }
}
