Add tests

Signed-off-by: Rishab Nahata <[email protected]>

Author: imRishN

server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java

@@ -52,6 +52,7 @@
 import org.opensearch.index.mapper.MappedFieldType;
 import org.opensearch.index.mapper.MapperService;
 import org.opensearch.index.mapper.NumberFieldMapper;
+import org.opensearch.index.query.BoolQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
 import org.opensearch.index.query.QueryShardContext;
 import org.opensearch.index.query.RangeQueryBuilder;
@@ -245,7 +246,9 @@ private void testStarTreeDocValuesInternal(Codec codec, List<DimensionFieldData>
         for (int cases = 0; cases < 15; cases++) {
             // Get all types of queries (Term/Terms/Range) for all the given dimensions.
             List<QueryBuilder> allFieldQueries = dimensionFieldData.stream()
-                .flatMap(x -> Stream.of(x.getTermQueryBuilder(), x.getTermsQueryBuilder(), x.getRangeQueryBuilder()))
+                .flatMap(
+                    x -> Stream.of(x.getTermQueryBuilder(), x.getTermsQueryBuilder(), x.getRangeQueryBuilder(), x.getBoolQueryBuilder())
+                )
                 .toList();
 
             for (QueryBuilder qb : allFieldQueries) {
@@ -547,6 +550,24 @@ public QueryBuilder getRangeQueryBuilder() {
                 .includeUpper(randomBoolean());
         }
 
+        public QueryBuilder getBoolQueryBuilder() {
+            // MUST only
+            BoolQueryBuilder mustOnly = new BoolQueryBuilder().must(getTermQueryBuilder()).must(getRangeQueryBuilder());
+
+            // MUST with nested SHOULD on same dimension
+            BoolQueryBuilder mustWithShould = new BoolQueryBuilder().must(getTermQueryBuilder())
+                .must(
+                    new BoolQueryBuilder().should(new TermQueryBuilder(fieldName, valueSupplier.get()))
+                        .should(new TermQueryBuilder(fieldName, valueSupplier.get()))
+                );
+
+            // SHOULD only on same dimension
+            BoolQueryBuilder shouldOnly = new BoolQueryBuilder().should(new TermQueryBuilder(fieldName, valueSupplier.get()))
+                .should(new RangeQueryBuilder(fieldName).from(valueSupplier.get()).to(valueSupplier.get()));
+
+            return randomFrom(mustOnly, mustWithShould, shouldOnly);
+        }
+
         public String getFieldType() {
             return fieldType;
         }

server/src/test/java/org/opensearch/search/aggregations/startree/StarTreeFilterTests.java

@@ -500,21 +500,17 @@ public void testStarTreeFilterWithBoolQueries() throws IOException {
             .thenReturn(new NumberFieldMapper.NumberFieldType(DV, NumberFieldMapper.NumberType.INTEGER));
 
         // Test 'MUST' clause
-        StarTreeFilter mustFilter = new StarTreeFilter(Map.of(
-            SNDV, List.of(new ExactMatchDimFilter(SNDV, List.of(0L))),
-            DV, List.of(new ExactMatchDimFilter(DV, List.of(0L)))
-        ));
+        StarTreeFilter mustFilter = new StarTreeFilter(
+            Map.of(SNDV, List.of(new ExactMatchDimFilter(SNDV, List.of(0L))), DV, List.of(new ExactMatchDimFilter(DV, List.of(0L))))
+        );
         long starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, mustFilter, context, searchContext);
         long docCount = getDocCount(docs, Map.of(SNDV, 0L, DV, 0L));
         assertEquals(docCount, starTreeDocCount);
 
         // Test 'SHOULD' clause (same dimension)
-        StarTreeFilter shouldFilter = new StarTreeFilter(Map.of(
-            SNDV, Arrays.asList(
-                new ExactMatchDimFilter(SNDV, List.of(0L)),
-                new ExactMatchDimFilter(SNDV, List.of(1L))
-            )
-        ));
+        StarTreeFilter shouldFilter = new StarTreeFilter(
+            Map.of(SNDV, Arrays.asList(new ExactMatchDimFilter(SNDV, List.of(0L)), new ExactMatchDimFilter(SNDV, List.of(1L))))
+        );
         starTreeDocCount = getDocCountFromStarTree(starTreeDocValuesReader, shouldFilter, context, searchContext);
         docCount = getDocCount(docs, Map.of(SNDV, 0L)) + getDocCount(docs, Map.of(SNDV, 1L));
         assertEquals(docCount, starTreeDocCount);