Fix unstable integration tests

integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/ppl/FieldsCommandIT.java

@@ -15,12 +15,6 @@
 
 package com.amazon.opendistroforelasticsearch.sql.ppl;
 
-import org.json.JSONObject;
-import org.junit.Ignore;
-import org.junit.jupiter.api.Test;
-
-import java.io.IOException;
-
 import static com.amazon.opendistroforelasticsearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
 import static com.amazon.opendistroforelasticsearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
 import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.columnName;
@@ -31,6 +25,11 @@
 import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.verifyDataRows;
 import static com.amazon.opendistroforelasticsearch.sql.util.MatcherUtils.verifySchema;
 
+import java.io.IOException;
+import org.json.JSONObject;
+import org.junit.Ignore;
+import org.junit.jupiter.api.Test;
+
 public class FieldsCommandIT extends PPLIntegTestCase {
 
   @Override
@@ -63,43 +62,18 @@ public void testFieldsWildCard() throws IOException {
 
   @Test
   public void testSelectDateTypeField() throws IOException {
-    String result =
-            executeQueryToString(
-                    String.format("source=%s | fields birthdate", TEST_INDEX_BANK));
-    assertEquals(
-        "{\n"
-            + "  \"schema\": [\n"
-            + "    {\n"
-            + "      \"name\": \"birthdate\",\n"
-            + "      \"type\": \"timestamp\"\n"
-            + "    }\n"
-            + "  ],\n"
-            + "  \"datarows\": [\n"
-            + "    [\n"
-            + "      \"2017-10-23 00:00:00\"\n"
-            + "    ],\n"
-            + "    [\n"
-            + "      \"2017-11-20 00:00:00\"\n"
-            + "    ],\n"
-            + "    [\n"
-            + "      \"2018-06-23 00:00:00\"\n"
-            + "    ],\n"
-            + "    [\n"
-            + "      \"2018-11-13 23:33:20\"\n"
-            + "    ],\n"
-            + "    [\n"
-            + "      \"2018-06-27 00:00:00\"\n"
-            + "    ],\n"
-            + "    [\n"
-            + "      \"2018-08-19 00:00:00\"\n"
-            + "    ],\n"
-            + "    [\n"
-            + "      \"2018-08-11 00:00:00\"\n"
-            + "    ]\n"
-            + "  ],\n"
-            + "  \"total\": 7,\n"
-            + "  \"size\": 7\n"
-            + "}\n",
-        result);
+    JSONObject result =
+        executeQuery(String.format("source=%s | fields birthdate", TEST_INDEX_BANK));
+    verifySchema(result, schema("birthdate", null, "timestamp"));
+
+    verifyDataRows(result,
+        rows("2017-10-23 00:00:00"),
+        rows("2017-11-20 00:00:00"),
+        rows("2018-06-23 00:00:00"),
+        rows("2018-11-13 23:33:20"),
+        rows("2018-06-27 00:00:00"),
+        rows("2018-08-19 00:00:00"),
+        rows("2018-08-11 00:00:00")
+    );
   }
 }

integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/ppl/StandaloneIT.java

@@ -35,11 +35,13 @@
 import com.amazon.opendistroforelasticsearch.sql.storage.StorageEngine;
 import com.google.common.collect.ImmutableMap;
 import java.io.IOException;
+import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
 import org.elasticsearch.client.Node;
 import org.elasticsearch.client.Request;
 import org.elasticsearch.client.RestClient;
+import org.elasticsearch.client.RestClientBuilder;
 import org.elasticsearch.client.RestHighLevelClient;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
@@ -58,8 +60,8 @@ public class StandaloneIT extends PPLIntegTestCase {
 
   @Override
   public void init() {
-    restClient =
-        new RestHighLevelClient(RestClient.builder(client().getNodes().toArray(new Node[0])));
+    // Using client() defined in ODFERestTestCase.
+    restClient = new InternalRestHighLevelClient(client());
 
     ElasticsearchClient client = new ElasticsearchRestClient(restClient);
     AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext();
@@ -73,12 +75,6 @@ public void init() {
     pplService = context.getBean(PPLService.class);
   }
 
-  @AfterEach
-  public void tearDown() throws Exception {
-    restClient.close();
-    super.tearDown();
-  }
-
   @Test
   public void testSourceFieldQuery() throws IOException {
     Request request1 = new Request("PUT", "/test/_doc/1?refresh=true");
@@ -144,4 +140,13 @@ public <T> T getSettingValue(Key key) {
       }
     };
   }
+
+  /**
+   * Internal RestHighLevelClient only for testing purpose.
+   */
+  static class InternalRestHighLevelClient extends RestHighLevelClient {
+    public InternalRestHighLevelClient(RestClient restClient) {
+      super(restClient, RestClient::close, Collections.emptyList());
+    }
+  }
 }

integ-test/src/test/java/com/amazon/opendistroforelasticsearch/sql/util/MatcherUtils.java

@@ -35,6 +35,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.function.Function;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.SearchHits;
 import org.hamcrest.Description;
@@ -46,6 +48,8 @@
 
 public class MatcherUtils {
 
+  private static final Logger LOG = LogManager.getLogger();
+
   /**
    * Assert field value in object by a custom matcher and getter to access the field.
    *
@@ -135,7 +139,12 @@ public static Matcher<JSONObject> kvInt(String key, Matcher<Integer> matcher) {
 
   @SafeVarargs
   public static void verifySchema(JSONObject response, Matcher<JSONObject>... matchers) {
-    verify(response.getJSONArray("schema"), matchers);
+    try {
+      verify(response.getJSONArray("schema"), matchers);
+    } catch (Exception e) {
+      LOG.error(String.format("verify schema failed, response: %s", response.toString()), e);
+      throw e;
+    }
   }
 
   @SafeVarargs