From c3d0da3204f2a46f09d8415d31f985dd65e73c16 Mon Sep 17 00:00:00 2001
From: Dai <daichen@amazon.com>
Date: Tue, 15 Dec 2020 14:07:49 -0800
Subject: [PATCH] Fix UT

---
 .../protocol/response/QueryResultTest.java    | 22 +++++++++++++++----
 .../format/CsvResponseFormatterTest.java      |  8 +++----
 .../SimpleJsonResponseFormatterTest.java      | 19 ++++++++++++++--
 3 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java
index ac17f90391..5373650b25 100644
--- a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java
+++ b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java
@@ -33,7 +33,7 @@
 class QueryResultTest {
 
   private ExecutionEngine.Schema schema = new ExecutionEngine.Schema(ImmutableList.of(
-      new ExecutionEngine.Schema.Column("name", "n", STRING),
+      new ExecutionEngine.Schema.Column("name", null, STRING),
       new ExecutionEngine.Schema.Column("age", null, INTEGER)));
 
 
@@ -58,7 +58,21 @@ void columnNameTypes() {
         ));
 
     assertEquals(
-        ImmutableMap.of("n", "string", "age", "integer"),
+        ImmutableMap.of("name", "string", "age", "integer"),
+        response.columnNameTypes()
+    );
+  }
+
+  @Test
+  void columnNameTypesWithAlias() {
+    ExecutionEngine.Schema schema = new ExecutionEngine.Schema(ImmutableList.of(
+        new ExecutionEngine.Schema.Column("name", "n", STRING)));
+    QueryResult response = new QueryResult(
+        schema,
+        Collections.singletonList(tupleValue(ImmutableMap.of("n", "John"))));
+
+    assertEquals(
+        ImmutableMap.of("n", "string"),
         response.columnNameTypes()
     );
   }
@@ -69,7 +83,7 @@ void columnNameTypesFromEmptyExprValues() {
         schema,
         Collections.emptyList());
     assertEquals(
-        ImmutableMap.of("n", "string", "age", "integer"),
+        ImmutableMap.of("name", "string", "age", "integer"),
         response.columnNameTypes()
     );
   }
@@ -84,7 +98,7 @@ void columnNameTypesFromExprValuesWithMissing() {
         ));
 
     assertEquals(
-        ImmutableMap.of("n", "string", "age", "integer"),
+        ImmutableMap.of("name", "string", "age", "integer"),
         response.columnNameTypes()
     );
   }
diff --git a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/CsvResponseFormatterTest.java b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/CsvResponseFormatterTest.java
index fdae598121..f720a3e775 100644
--- a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/CsvResponseFormatterTest.java
+++ b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/CsvResponseFormatterTest.java
@@ -53,10 +53,10 @@ void formatResponse() {
   @Test
   void sanitizeHeaders() {
     ExecutionEngine.Schema schema = new ExecutionEngine.Schema(ImmutableList.of(
-        new ExecutionEngine.Schema.Column("=firstname", "firstname", STRING),
-        new ExecutionEngine.Schema.Column("+lastname", "lastname", STRING),
-        new ExecutionEngine.Schema.Column("-city", "city", STRING),
-        new ExecutionEngine.Schema.Column("@age", "age", INTEGER)));
+        new ExecutionEngine.Schema.Column("=firstname", null, STRING),
+        new ExecutionEngine.Schema.Column("+lastname", null, STRING),
+        new ExecutionEngine.Schema.Column("-city", null, STRING),
+        new ExecutionEngine.Schema.Column("@age", null, INTEGER)));
     QueryResult response = new QueryResult(schema, Arrays.asList(
         tupleValue(ImmutableMap.of(
             "=firstname", "John", "+lastname", "Smith", "-city", "Seattle", "@age", 20))));
diff --git a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java
index f7d8d6e710..6c72ed8d38 100644
--- a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java
+++ b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java
@@ -36,8 +36,8 @@
 class SimpleJsonResponseFormatterTest {
 
   private final ExecutionEngine.Schema schema = new ExecutionEngine.Schema(ImmutableList.of(
-      new ExecutionEngine.Schema.Column("firstname", "name", STRING),
-      new ExecutionEngine.Schema.Column("age", "age", INTEGER)));
+      new ExecutionEngine.Schema.Column("firstname", null, STRING),
+      new ExecutionEngine.Schema.Column("age", null, INTEGER)));
 
   @Test
   void formatResponse() {
@@ -92,6 +92,21 @@ void formatResponsePretty() {
         formatter.format(response));
   }
 
+  @Test
+  void formatResponseSchemaWithAlias() {
+    ExecutionEngine.Schema schema = new ExecutionEngine.Schema(ImmutableList.of(
+        new ExecutionEngine.Schema.Column("firstname", "name", STRING)));
+    QueryResult response =
+        new QueryResult(
+            schema,
+            ImmutableList.of(tupleValue(ImmutableMap.of("name", "John", "age", 20))));
+    SimpleJsonResponseFormatter formatter = new SimpleJsonResponseFormatter(COMPACT);
+    assertEquals(
+        "{\"schema\":[{\"name\":\"name\",\"type\":\"string\"}],"
+            + "\"datarows\":[[\"John\",20]],\"total\":1,\"size\":1}",
+        formatter.format(response));
+  }
+
   @Test
   void formatResponseWithMissingValue() {
     QueryResult response =