/*
 * Copyright (c) 2014 Spotify AB.
 *
 * 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.spotify.helios.testing;

import com.google.common.net.HostAndPort;
import com.google.common.util.concurrent.ListenableFuture;

import com.spotify.helios.client.HeliosClient;
import com.spotify.helios.common.descriptors.HostStatus;

import org.junit.rules.ExternalResource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.ConnectException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketAddress;
import java.net.SocketTimeoutException;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * A HeliosDeploymentResource makes the supplied {@link HeliosDeployment} available to a JUnit
 * test, and guarantees to tear it down afterward.
 */
public class HeliosDeploymentResource extends ExternalResource {

  private static final Logger log = LoggerFactory.getLogger(HeliosDeploymentResource.class);

  private final HeliosDeployment deployment;

  /**
   * @param deployment The Helios deployment to expose to your JUnit tests.
   */
  public HeliosDeploymentResource(final HeliosDeployment deployment) {
    this.deployment = deployment;
  }

  /** Ensure that the HeliosDeployment is up. */
  @Override
  public void before() throws Throwable {
    super.before();

    // wait for the helios master to be available
    Polling.awaitUnchecked(30, TimeUnit.SECONDS, new Callable<Boolean>() {
      @Override
      public Boolean call() throws Exception {
        final HostAndPort hap = deployment.address();
        final SocketAddress address = new InetSocketAddress(hap.getHostText(), hap.getPort());
        log.debug("attempting to connect to {}", address);

        try {
          final Socket s = new Socket();
          s.connect(address, 100);
          log.info("successfully connected to address {} for {}", address, deployment);
          return true;
        } catch (SocketTimeoutException | ConnectException e) {
          log.debug("could not yet connect to HeliosDeployment: {}", e.toString());
          return null;
        }
      }
    });

    // Ensure that at least one agent is available and UP in this HeliosDeployment.
    // This prevents continuing with the test when starting up helios-solo before the agent is
    // registered.
    final HeliosClient client = client();
    Polling.awaitUnchecked(30, TimeUnit.SECONDS, new Callable<Boolean>() {
      @Override
      public Boolean call() throws Exception {
        final ListenableFuture<List<String>> future = client.listHosts();

        final List<String> hosts;
        try {
          // use a short timeout to allow this request to be retried a few times by the
          // Polling.await loop
           hosts = future.get(1, TimeUnit.SECONDS);
        } catch (TimeoutException | InterruptedException e) {
          log.debug("timed out waiting for listHosts request to finish, will retry");
          return null;
        }

        if (hosts.isEmpty()) {
          log.debug("0 agents in {}, will retry", deployment);
          return null;
        }

        log.info("Ensured that at least one agent has registered with ZooKeeper in this"
                 + "HeliosDeployment. Now checking if it is UP.");

        // Check that at least one host is UP (is maintaining a reasonably reliable
        // connection to ZK) in addition to registering.
        final ListenableFuture<HostStatus> statusFuture = client.hostStatus(hosts.get(0));
        final HostStatus hostStatus;
        try {
          hostStatus = statusFuture.get(1, TimeUnit.SECONDS);
        } catch (TimeoutException | InterruptedException e) {
          log.debug("timed out waiting for hostStatus request on host {} to finish, will retry",
                    hosts.get(0));
          return null;
        }

        if (hostStatus != null && hostStatus.getStatus() == HostStatus.Status.UP) {
          log.info("Ensured that at least one agent is UP in this HeliosDeployment, "
                   + "continuing with test!");
          return true;
        }

        return null;
      }
    });
  }

  @Override
  public void after() {
    log.info("Tearing down {}", this.deployment);
    deployment.close();
  }

  /**
   *
   * @return A Helios client connected to the Helios deployment supplied when instantiating the
   * HeliosDeploymentResource.
   */
  public HeliosClient client() {
    return deployment.client();
  }

}

@@ -1125,16 +1125,38 @@
 l.List;%0A
+import java.util.Map;%0A
 import j
@@ -3735,165 +3735,8 @@
  %7D%0A%0A
-        log.info(%22Ensured that at least one agent has registered with ZooKeeper in this%22%0A                 + %22HeliosDeployment. Now checking if it is UP.%22);%0A%0A
     
@@ -3903,16 +3903,28 @@
 eFuture%3C
+Map%3CString, 
 HostStat
@@ -3926,16 +3926,17 @@
 tStatus%3E
+%3E
  statusF
@@ -3964,21 +3964,16 @@
 atus
+es
 (hosts
-.get(0)
 );%0A 
@@ -3980,24 +3980,36 @@
        final
+ Map%3CString,
  HostStatus 
@@ -4007,16 +4007,17 @@
 stStatus
+%3E
  hostSta
@@ -4019,16 +4019,18 @@
 stStatus
+es
 ;%0A      
@@ -4057,16 +4057,18 @@
 stStatus
+es
  = statu
@@ -4220,27 +4220,10 @@
 atus
- request on host %7B%7D
+es
  to 
@@ -4241,17 +4241,18 @@
 l retry%22
-,
+);
 %0A       
@@ -4258,67 +4258,102 @@
     
-          hosts.get(0));%0A          return null;%0A        %7D%0A%0A
+return null;%0A        %7D%0A%0A        for (final HostStatus hostStatus : hostStatuses.values()) %7B%0A  
     
@@ -4434,32 +4434,34 @@
 UP) %7B%0A          
+  
 log.info(%22Ensure
@@ -4539,16 +4539,18 @@
         
+ 
+ 
 + %22conti
@@ -4571,32 +4571,34 @@
 t!%22);%0A          
+  
 return true;%0A   
@@ -4586,32 +4586,44 @@
    return true;%0A
+          %7D%0A
         %7D%0A%0A     

Address PR comments