add tests

Signed-off-by: bowenlan-amzn <[email protected]>

Author: bowenlan-amzn

server/src/main/java/org/opensearch/common/settings/FeatureFlagSettings.java

@@ -39,6 +39,7 @@ protected FeatureFlagSettings(
         FeatureFlags.TERM_VERSION_PRECOMMIT_ENABLE_SETTING,
         FeatureFlags.ARROW_STREAMS_SETTING,
         FeatureFlags.STREAM_TRANSPORT_SETTING,
-        FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING
+        FeatureFlags.MERGED_SEGMENT_WARMER_EXPERIMENTAL_SETTING,
+        FeatureFlags.STREAM_SEARCH_SETTING
     );
 }

server/src/main/java/org/opensearch/rest/action/search/RestSearchAction.java

@@ -446,12 +446,11 @@ public boolean allowsUnsafeBuffers() {
 
     /**
      * Determines if a search request can use stream search.
-     * Stream search is only supported for requests with exactly one terms aggregation and no other aggregations.
      *
      * @param searchRequest the search request to validate
      * @return true if the request can use stream search, false otherwise
      */
-    private static boolean canUseStreamSearch(SearchRequest searchRequest) {
+    static boolean canUseStreamSearch(SearchRequest searchRequest) {
         if (searchRequest.source() == null || searchRequest.source().aggregations() == null) {
             return true; // No aggregations, stream search is allowed
         }

server/src/test/java/org/opensearch/rest/action/search/RestSearchActionTests.java

@@ -11,12 +11,16 @@
 import org.opensearch.action.ActionRequest;
 import org.opensearch.action.ActionType;
 import org.opensearch.action.search.SearchAction;
+import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.StreamSearchAction;
 import org.opensearch.common.SetOnce;
 import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.core.action.ActionListener;
 import org.opensearch.core.action.ActionResponse;
+import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.rest.RestRequest;
+import org.opensearch.search.aggregations.AggregationBuilders;
+import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.tasks.Task;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.test.client.NoOpNodeClient;
@@ -91,4 +95,51 @@ public void testWithStreamSearchAndTransportEnabled() throws Exception {
             testActionExecution(StreamSearchAction.INSTANCE);
         }
     }
+
+    // Tests for canUseStreamSearch method
+    public void testCanUseStreamSearchWithNullSource() {
+        SearchRequest searchRequest = new SearchRequest();
+        assertTrue(RestSearchAction.canUseStreamSearch(searchRequest));
+    }
+
+    public void testCanUseStreamSearchWithNoAggregations() {
+        SearchRequest searchRequest = new SearchRequest();
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        source.query(QueryBuilders.matchAllQuery());
+        searchRequest.source(source);
+        assertTrue(RestSearchAction.canUseStreamSearch(searchRequest));
+    }
+
+    public void testCanUseStreamSearchWithSingleTermsAggregation() {
+        SearchRequest searchRequest = new SearchRequest();
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        source.aggregation(AggregationBuilders.terms("test_terms").field("category"));
+        searchRequest.source(source);
+        assertTrue(RestSearchAction.canUseStreamSearch(searchRequest));
+    }
+
+    public void testCanUseStreamSearchWithMultipleAggregations() {
+        SearchRequest searchRequest = new SearchRequest();
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        source.aggregation(AggregationBuilders.terms("test_terms").field("category"));
+        source.aggregation(AggregationBuilders.avg("test_avg").field("price"));
+        searchRequest.source(source);
+        assertFalse(RestSearchAction.canUseStreamSearch(searchRequest));
+    }
+
+    public void testCanUseStreamSearchWithSingleNonTermsAggregation() {
+        SearchRequest searchRequest = new SearchRequest();
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        source.aggregation(AggregationBuilders.avg("test_avg").field("price"));
+        searchRequest.source(source);
+        assertFalse(RestSearchAction.canUseStreamSearch(searchRequest));
+    }
+
+    public void testCanUseStreamSearchWithSingleHistogramAggregation() {
+        SearchRequest searchRequest = new SearchRequest();
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        source.aggregation(AggregationBuilders.histogram("test_histogram").field("timestamp").interval(1000));
+        searchRequest.source(source);
+        assertFalse(RestSearchAction.canUseStreamSearch(searchRequest));
+    }
 }