Added more cases to improve coverage

Signed-off-by: expani <[email protected]>
opensearch-project · Jan 3, 2025 · 6bf70e4 · 6bf70e4
1 parent 6432da4
commit 6bf70e4
Showing 1 changed file with 53 additions and 6 deletions.
diff --git a/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java b/server/src/test/java/org/opensearch/search/aggregations/startree/MetricAggregatorTests.java
@@ -64,6 +64,7 @@
 import org.opensearch.search.aggregations.metrics.ValueCountAggregationBuilder;
 import org.junit.After;
 import org.junit.Before;
+import org.opensearch.search.aggregations.support.ValuesSourceAggregatorFactory;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -74,8 +75,8 @@
 import java.util.function.BiConsumer;
 import java.util.function.Function;
 
-import org.mockito.Mockito;
-
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 import static org.opensearch.search.aggregations.AggregationBuilders.avg;
 import static org.opensearch.search.aggregations.AggregationBuilders.count;
 import static org.opensearch.search.aggregations.AggregationBuilders.max;
@@ -290,10 +291,10 @@ public void testStarTreeDocValues() throws IOException {
             new MockBigArrays(new MockPageCacheRecycler(Settings.EMPTY), circuitBreakerService).withCircuitBreaking()
         );
 
-        MetricAggregatorFactory aggregatorFactory = Mockito.mock(MetricAggregatorFactory.class);
-        Mockito.when(aggregatorFactory.getSubFactories()).thenReturn(AggregatorFactories.EMPTY);
-        Mockito.when(aggregatorFactory.getField()).thenReturn(FIELD_NAME);
-        Mockito.when(aggregatorFactory.getMetricStat()).thenReturn(MetricStat.SUM);
+        MetricAggregatorFactory aggregatorFactory = mock(MetricAggregatorFactory.class);
+        when(aggregatorFactory.getSubFactories()).thenReturn(AggregatorFactories.EMPTY);
+        when(aggregatorFactory.getField()).thenReturn(FIELD_NAME);
+        when(aggregatorFactory.getMetricStat()).thenReturn(MetricStat.SUM);
 
         // Case when field and metric type in aggregation are fully supported by star tree.
         testCase(
@@ -338,6 +339,52 @@ public void testStarTreeDocValues() throws IOException {
             false
         );
 
+        // Case when field is not present in supported metrics
+        testCase(
+            indexSearcher,
+            query,
+            queryBuilder,
+            sumAggregationBuilder,
+            starTree,
+            supportedDimensions,
+            List.of(new Metric("hello", List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
+            verifyAggregation(InternalSum::getValue),
+            aggregatorFactory,
+            false
+        );
+
+        AggregatorFactories aggregatorFactories = mock(AggregatorFactories.class);
+        when(aggregatorFactories.getFactories()).thenReturn(new AggregatorFactory[]{mock(MetricAggregatorFactory.class)});
+        when(aggregatorFactory.getSubFactories()).thenReturn(aggregatorFactories);
+
+        // Case when sub aggregations are present
+        testCase(
+            indexSearcher,
+            query,
+            queryBuilder,
+            sumAggregationBuilder,
+            starTree,
+            supportedDimensions,
+            List.of(new Metric("hello", List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
+            verifyAggregation(InternalSum::getValue),
+            aggregatorFactory,
+            false
+        );
+
+        // Case when aggregation factory is not metric aggregation
+        testCase(
+            indexSearcher,
+            query,
+            queryBuilder,
+            sumAggregationBuilder,
+            starTree,
+            supportedDimensions,
+            List.of(new Metric("hello", List.of(MetricStat.MAX, MetricStat.MIN, MetricStat.AVG))),
+            verifyAggregation(InternalSum::getValue),
+            mock(ValuesSourceAggregatorFactory.class),
+            false
+        );
+
         ir.close();
         directory.close();
     }