Fix UT

opendistro-for-elasticsearch · Dec 15, 2020 · c3d0da3 · c3d0da3
1 parent 72d2a94
commit c3d0da3
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 10 deletions.
diff --git a/...est/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResultTest.java b/...est/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/...zon/opendistroforelasticsearch/sql/protocol/response/format/CsvResponseFormatterTest.java b/...zon/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/...ndistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatterTest.java b/...ndistroforelasticsearch/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 =