Adding test to cover tuple value extraction in HighlightOperator.

Signed-off-by: forestmvey <forestv@bitquilltech.com>
Bit-Quill · Sep 1, 2022 · ef316d6 · ef316d6
1 parent 986bb99
commit ef316d6
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/planner/physical/HighlightOperator.java b/core/src/main/java/org/opensearch/sql/planner/physical/HighlightOperator.java
@@ -7,7 +7,6 @@
 
 import static org.opensearch.sql.data.type.ExprCoreType.STRING;
 import static org.opensearch.sql.data.type.ExprCoreType.STRUCT;
-import static org.opensearch.sql.data.type.ExprCoreType.UNDEFINED;
 import static org.opensearch.sql.expression.env.Environment.extendEnv;
 
 import com.google.common.collect.ImmutableMap;
@@ -89,7 +88,7 @@ private Pair<String, ExprValue> mapHighlight(Environment<Expression, ExprValue>
     // used in conjunction with other highlight calls, we need to ensure
     // only wildcard regex matching is mapped to wildcard call.
     if (StringUtils.unquoteText(highlight.toString()).matches("(.+\\*)|(\\*.+)")
-        && value.type() != UNDEFINED) {
+        && value.type() == STRUCT) {
       value = new ExprTupleValue(
           new LinkedHashMap<String, ExprValue>(value.tupleValue()
               .entrySet()

diff --git a/core/src/test/java/org/opensearch/sql/planner/physical/HighlightOperatorTest.java b/core/src/test/java/org/opensearch/sql/planner/physical/HighlightOperatorTest.java
@@ -21,6 +21,7 @@
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.opensearch.sql.data.model.ExprNullValue;
 import org.opensearch.sql.data.model.ExprValue;
 import org.opensearch.sql.data.model.ExprValueUtils;
 import org.opensearch.sql.expression.DSL;
@@ -75,8 +76,12 @@ public void highlight_one_field() {
 
   @Test
   public void highlight_wildcard() {
-    when(inputPlan.hasNext()).thenReturn(true,false);
+    when(inputPlan.hasNext()).thenReturn(true, true, false);
     when(inputPlan.next())
+        .thenReturn(
+            tupleValue(ImmutableMap.of(
+                "_highlight", ExprNullValue.of(),
+                "action", "GET", "response", 200)))
         .thenReturn(
             tupleValue(ImmutableMap.of(
                 "_highlight", tupleValue(
@@ -86,13 +91,19 @@ public void highlight_wildcard() {
     assertThat(
         execute(new HighlightOperator(inputPlan, DSL.ref("r*", STRING))),
         contains(
+            tupleValue(ImmutableMap.of(
+                "_highlight", ExprNullValue.of(),
+                "action", "GET",
+                "response", 200, "highlight(r*)", ExprNullValue.of())
+            ),
             tupleValue(ImmutableMap.of(
                 "_highlight", tupleValue(
                       ImmutableMap.of("region", "us-east-1", "country", "us")),
                 "action", "GET",
                 "response", 200, "highlight(r*)", tupleValue(
                       ImmutableMap.of("region", "us-east-1")))
-        ))
+            )
+        )
     );
   }
-}
+}