Skip to content

Commit

Permalink
Added QueryShapeServiceTests
Browse files Browse the repository at this point in the history
Signed-off-by: David Zane <davizane@amazon.com>
  • Loading branch information
dzane17 committed Jul 24, 2024
1 parent 2f995db commit 38eef6a
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
Original file line number Diff line number Diff line change
@@ -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);
}
}

0 comments on commit 38eef6a

Please sign in to comment.