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.containers.DataPusherContainer;
import org.json.JSONArray;
import org.json.JSONObject;
import org.junit.BeforeClass;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TestWatcher;
import org.junit.runner.Description;

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


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

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

    private static DataPusherContainer pusher;

    private static List<String> slavesElasticAddresses = new ArrayList<>();

    @BeforeClass
    public static void startDataPusher() {

        try {
            List<JSONObject> tasks = new TasksResponse(new ESTasks(TEST_CONFIG, getScheduler().getIpAddress()), CLUSTER.getConfig().getNumberOfSlaves()).getTasks();
            for (JSONObject task : tasks) {
                LOGGER.info(task);
                slavesElasticAddresses.add(task.getString("http_address"));

            }
        } catch (Exception e) {
            LOGGER.error("Exception thrown: " + e.getMessage());
            throw new RuntimeException(e.getMessage());
        }

        Awaitility.await().atMost(5, TimeUnit.MINUTES).pollDelay(2, TimeUnit.SECONDS).until(() -> {
            try {
                if (!(Unirest.get("http://" + getSlavesElasticAddresses().get(0) + "/_nodes").asJson().getBody().getObject().getJSONObject("nodes").length() == 3)) {
                    return false;
                }
                pusher = new DataPusherContainer(CLUSTER.getConfig().dockerClient, getSlavesElasticAddresses().get(0));
                CLUSTER.addAndStartContainer(pusher);
                return true;
            } catch (UnirestException e) {
                return false;
            }
        });
    }

    @Test
    public void testDataConsistency() throws Exception {
        LOGGER.info("Addresses:");
        LOGGER.info(getSlavesElasticAddresses());
        Awaitility.await().atMost(1, TimeUnit.MINUTES).pollDelay(2, TimeUnit.SECONDS).until(() -> {
            JSONArray responseElements;
            for (String httpAddress: DataRetrievableAllNodesSystemTest.getSlavesElasticAddresses()) {
                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;
        });
    }

    public static DataPusherContainer getPusher() {
        return pusher;
    }

    public static List<String> getSlavesElasticAddresses() {
        return slavesElasticAddresses;
    }

}
@@ -293,16 +293,103 @@
 stBase;%0A
+import org.apache.mesos.elasticsearch.systemtest.callbacks.ElasticsearchNodesResponse;%0A
 import o
@@ -543,132 +543,31 @@
 fore
-Class;%0Aimport org.junit.Rule;%0Aimport org.junit.Test;%0Aimport org.junit.rules.TestWatcher;%0Aimport org.junit.runner.Description
+;%0Aimport org.junit.Test
 ;%0A%0Ai
@@ -1005,13 +1005,8 @@
 fore
-Class
 %0A   
@@ -1013,23 +1013,16 @@
  public 
-static 
 void sta
@@ -1069,51 +1069,26 @@
     
-List%3CJSONObject%3E tasks = new TasksResponse(
+ESTasks esTasks = 
 new 
@@ -1142,62 +1142,319 @@
 s())
-, CLUSTER.getConfig().getNumberOfSlaves()).getTasks();
+;%0A            ElasticsearchNodesResponse elasticsearchNodesResponse = new ElasticsearchNodesResponse(esTasks, CLUSTER.getConfig().getNumberOfSlaves());%0A            if (!elasticsearchNodesResponse.isDiscoverySuccessful()) %7B%0A                throw new RuntimeException(%22Could not discover ES nodes%22);%0A            %7D
 %0A   
@@ -1489,13 +1489,26 @@
 k : 
-t
+esTasks.getT
 asks
+()
 ) %7B%0A
@@ -1618,17 +1618,16 @@
 ess%22));%0A
-%0A
         
@@ -1709,33 +1709,19 @@
 hrown: %22
- + e.getMessage()
+, e
 );%0A     
@@ -3365,90 +3365,8 @@
  %7D%0A%0A
-    public static DataPusherContainer getPusher() %7B%0A        return pusher;%0A    %7D%0A%0A
     

Make data pusher run in before, not before class. May run too early. Some refactoring.