Added QueryShapeServiceTests

Signed-off-by: David Zane <[email protected]>

src/main/java/org/opensearch/plugin/insights/core/service/categorizer/QueryShapeService.java

@@ -22,10 +22,10 @@
 import java.util.Collections;
 import java.util.List;
 
-class QueryShapeService {
+public class QueryShapeService {
     static final String TWO_SPACE_INDENT = "  ";
 
-    static String buildShape(SearchSourceBuilder source, Boolean showFields) {
+    public static String buildShape(SearchSourceBuilder source, Boolean showFields) {
         StringBuilder shape = new StringBuilder();
         shape.append(buildQueryShape(source.query()));
         shape.append(buildAggregationShape(source.aggregations(), showFields));

src/test/java/org/opensearch/plugin/insights/core/service/categorizor/QueryShapeServiceTests.java

@@ -0,0 +1,73 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.plugin.insights.core.service.categorizor;
+
+import org.opensearch.index.query.BoolQueryBuilder;
+import org.opensearch.index.query.MatchQueryBuilder;
+import org.opensearch.index.query.QueryBuilders;
+import org.opensearch.index.query.RegexpQueryBuilder;
+import org.opensearch.index.query.TermQueryBuilder;
+import org.opensearch.plugin.insights.core.service.categorizer.QueryShapeService;
+import org.opensearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
+import org.opensearch.search.aggregations.support.ValueType;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.search.sort.SortOrder;
+import org.opensearch.test.OpenSearchTestCase;
+
+public final class QueryShapeServiceTests extends OpenSearchTestCase {
+    public void testQueryShape() {
+        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+        sourceBuilder.size(0);
+        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery("field", "value2");
+        MatchQueryBuilder matchQueryBuilder = QueryBuilders.matchQuery("tags", "php");
+        RegexpQueryBuilder regexpQueryBuilder = new RegexpQueryBuilder("field", "text");
+        sourceBuilder.query(new BoolQueryBuilder().must(termQueryBuilder).filter(matchQueryBuilder).should(regexpQueryBuilder));
+
+        String shape = QueryShapeService.buildShape(sourceBuilder, true);
+
+        String expected = "bool\n" + "  must:\n" + "    term\n" + "  filter:\n" + "    match\n" + "  should:\n" + "    regexp\n";
+        assertEquals(expected, shape);
+    }
+
+    public void testAggregationShape() {
+        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+        sourceBuilder.aggregation(new TermsAggregationBuilder("agg1").userValueTypeHint(ValueType.STRING).field("type"));
+        sourceBuilder.aggregation(new TermsAggregationBuilder("agg2").userValueTypeHint(ValueType.STRING).field("model"));
+        sourceBuilder.aggregation(
+            new TermsAggregationBuilder("agg3").userValueTypeHint(ValueType.STRING)
+                .field("key")
+                .subAggregation(new TermsAggregationBuilder("child-agg1").userValueTypeHint(ValueType.STRING).field("key.sub1"))
+                .subAggregation(new TermsAggregationBuilder("child-agg2").userValueTypeHint(ValueType.STRING).field("key.sub2"))
+        );
+
+        String shape = QueryShapeService.buildShape(sourceBuilder, true);
+
+        String expected = "aggregation:\n"
+            + "  terms [key]\n"
+            + "    aggregation:\n"
+            + "      terms [key.sub1]\n"
+            + "      terms [key.sub2]\n"
+            + "  terms [model]\n"
+            + "  terms [type]\n";
+        assertEquals(expected, shape);
+    }
+
+    public void testSortShape() {
+        SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
+        sourceBuilder.sort("color", SortOrder.DESC);
+        sourceBuilder.sort("vendor", SortOrder.DESC);
+        sourceBuilder.sort("price", SortOrder.ASC);
+        sourceBuilder.sort("album", SortOrder.ASC);
+
+        String shape = QueryShapeService.buildShape(sourceBuilder, true);
+
+        String expected = "sort:\n" + "  asc [album]\n" + "  asc [price]\n" + "  desc [color]\n" + "  desc [vendor]\n";
+        assertEquals(expected, shape);
+    }
+}