package org.apache.mesos.elasticsearch.systemtest;

import com.jayway.awaitility.Awaitility;
import com.mashape.unirest.http.Unirest;
import com.mashape.unirest.http.exceptions.UnirestException;
import org.apache.log4j.Logger;
import org.apache.mesos.elasticsearch.systemtest.base.SchedulerTestBase;
import org.apache.mesos.elasticsearch.systemtest.callbacks.ElasticsearchNodesResponse;
import org.apache.mesos.elasticsearch.systemtest.containers.DataPusherContainer;
import org.json.JSONArray;
import org.json.JSONObject;
import org.junit.Before;
import org.junit.Test;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

import static org.junit.Assert.assertTrue;


/**
 * Tests scheduler APIs
 */
public class DataRetrievableAllNodesSystemTest extends SchedulerTestBase {

    private static final Logger LOGGER = Logger.getLogger(DataRetrievableAllNodesSystemTest.class);

    private static DataPusherContainer pusher;
    private List<String> esAddresses;

    @Test
    public void testDataConsistency() throws Exception {
        ESTasks esTasks = new ESTasks(TEST_CONFIG, getScheduler().getIpAddress());

        esAddresses = esTasks.getTasks().stream().map(task -> task.getString("http_address")).collect(Collectors.toList());

        Awaitility.await().atMost(5, TimeUnit.MINUTES).pollDelay(2, TimeUnit.SECONDS).until(() -> {
            try {
                // This may throw a JSONException if we call before the JSON has been generated. Hence, catch exception.
                if (!(Unirest.get("http://" + esAddresses.get(0) + "/_nodes").asJson().getBody().getObject().getJSONObject("nodes").length() == 3)) {
                    return false;
                }
                pusher = new DataPusherContainer(CLUSTER.getConfig().dockerClient, esAddresses.get(0));
                CLUSTER.addAndStartContainer(pusher);
                return true;
            } catch (Exception e) {
                return false;
            }
        });

        LOGGER.info("Addresses:");
        LOGGER.info(esAddresses);
        Awaitility.await().atMost(1, TimeUnit.MINUTES).pollDelay(2, TimeUnit.SECONDS).until(() -> {
            JSONArray responseElements;
            for (String httpAddress: esAddresses) {
                try {
                    responseElements = Unirest.get("http://" + httpAddress + "/_count").asJson().getBody().getArray();
                } catch (Exception e) {
                    LOGGER.error("Unirest exception:" + e.getMessage());
                    return false;
                }
                LOGGER.info(responseElements);
                if (9 > responseElements.getJSONObject(0).getInt("count")) {
                    return false;
                }
            }
            return true;
        });
    }
}
@@ -1185,251 +1185,257 @@
 ));%0A
-%0A        esAddresses = esTasks.getTasks().stream().map(task -%3E task.getString(%22http_address%22)).collect(Collectors.toList());%0A%0A        Awaitility.await().atMost(5, TimeUnit.MINUTES).pollDelay(2, TimeUnit.SECONDS).until(() -%3E %7B%0A            try %7B
+        Awaitility.await().atMost(5, TimeUnit.MINUTES).pollDelay(2, TimeUnit.SECONDS).until(() -%3E %7B%0A            try %7B%0A                esAddresses = esTasks.getTasks().stream().map(task -%3E task.getString(%22http_address%22)).collect(Collectors.toList());
 %0A   
@@ -1572,14 +1572,15 @@
     
-if (!(
+return 
 Unir
@@ -1698,25 +1698,54 @@
 h() == 3
-)) %7B
+;
 %0A    
+        %7D catch (Exception e) %7B%0A
         
@@ -1758,36 +1758,32 @@
   return false;%0A
-    
             %7D%0A  
@@ -1780,32 +1780,37 @@
       %7D%0A        
+%7D);%0A%0A
         pusher =
@@ -1885,32 +1885,24 @@
 es.get(0));%0A
-        
         CLUS
@@ -1938,129 +1938,8 @@
 er);
-%0A                return true;%0A            %7D catch (Exception e) %7B%0A                return false;%0A            %7D%0A        %7D);
 %0A%0A  

Move task refresh inside await. Move start container outside await.