Set APIs behind the feature flag

Signed-off-by: Owais Kazi <[email protected]>

src/main/java/org/opensearch/flowframework/rest/AbstractSearchWorkflowAction.java

@@ -13,6 +13,7 @@
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.flowframework.common.FlowFrameworkFeatureEnabledSetting;
 import org.opensearch.flowframework.exception.FlowFrameworkException;
@@ -23,19 +24,19 @@
 import org.opensearch.rest.RestResponse;
 import org.opensearch.rest.action.RestResponseListener;
 import org.opensearch.search.builder.SearchSourceBuilder;
-import static org.opensearch.flowframework.common.FlowFrameworkFeatureEnabledSetting.FLOW_FRAMEWORK_ENABLED;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
 
 import static org.opensearch.core.xcontent.ToXContent.EMPTY_PARAMS;
+import static org.opensearch.flowframework.common.FlowFrameworkFeatureEnabledSetting.FLOW_FRAMEWORK_ENABLED;
 import static org.opensearch.flowframework.util.RestHandlerUtils.getSourceContext;
 
 /**
  * Abstract class to handle search request.
  */
-public class AbstractSearchWorkflowAction<T extends ToXContentObject> extends BaseRestHandler {
+public abstract class AbstractSearchWorkflowAction<T extends ToXContentObject> extends BaseRestHandler {
 
     protected final List<String> urlPaths;
     protected final String index;
@@ -49,26 +50,31 @@ public class AbstractSearchWorkflowAction<T extends ToXContentObject> extends Ba
      * @param index index the search should be done on
      * @param clazz model class
      * @param actionType from which action abstract class is called
+     * @param flowFrameworkFeatureEnabledSetting Whether this API is enabled
      */
-    public AbstractSearchWorkflowAction(List<String> urlPaths, String index, Class<T> clazz, ActionType<SearchResponse> actionType, FlowFrameworkFeatureEnabledSetting flowFrameworkFeatureEnabledSetting) {
+    public AbstractSearchWorkflowAction(
+        List<String> urlPaths,
+        String index,
+        Class<T> clazz,
+        ActionType<SearchResponse> actionType,
+        FlowFrameworkFeatureEnabledSetting flowFrameworkFeatureEnabledSetting
+    ) {
         this.urlPaths = urlPaths;
         this.index = index;
         this.clazz = clazz;
         this.actionType = actionType;
         this.flowFrameworkFeatureEnabledSetting = flowFrameworkFeatureEnabledSetting;
     }
 
-    @Override
-    public String getName() {
-        return null;
-    }
-
     @Override
     protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
         if (!flowFrameworkFeatureEnabledSetting.isFlowFrameworkEnabled()) {
-            throw new FlowFrameworkException(
-                    "This API is disabled. To enable it, update the setting [" + FLOW_FRAMEWORK_ENABLED.getKey() + "] to true.",
-                    RestStatus.FORBIDDEN
+            FlowFrameworkException ffe = new FlowFrameworkException(
+                "This API is disabled. To enable it, update the setting [" + FLOW_FRAMEWORK_ENABLED.getKey() + "] to true.",
+                RestStatus.FORBIDDEN
+            );
+            return channel -> channel.sendResponse(
+                    new BytesRestResponse(ffe.getRestStatus(), ffe.toXContent(channel.newErrorBuilder(), ToXContent.EMPTY_PARAMS))
             );
         }
         SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();

src/main/java/org/opensearch/flowframework/rest/RestSearchWorkflowAction.java

@@ -25,8 +25,14 @@ public class RestSearchWorkflowAction extends AbstractSearchWorkflowAction<Templ
     private static final String SEARCH_WORKFLOW_PATH = WORKFLOW_URI + "/_search";
 
     public RestSearchWorkflowAction(FlowFrameworkFeatureEnabledSetting flowFrameworkFeatureEnabledSetting) {
-        super(ImmutableList.of(SEARCH_WORKFLOW_PATH), GLOBAL_CONTEXT_INDEX, Template.class, SearchWorkflowAction.INSTANCE, flowFrameworkFeatureEnabledSetting);
-
+        super(
+            ImmutableList.of(SEARCH_WORKFLOW_PATH),
+            GLOBAL_CONTEXT_INDEX,
+            Template.class,
+            SearchWorkflowAction.INSTANCE,
+            flowFrameworkFeatureEnabledSetting
+        );
+    }
 
     @Override
     public String getName() {

src/test/java/org/opensearch/flowframework/FlowFrameworkPluginTests.java

@@ -74,18 +74,13 @@ public void tearDown() throws Exception {
 
     public void testPlugin() throws IOException {
         try (FlowFrameworkPlugin ffp = new FlowFrameworkPlugin()) {
-<<<<<<< HEAD
             assertEquals(
                 3,
                 ffp.createComponents(client, clusterService, threadPool, null, null, null, environment, null, null, null, null).size()
             );
-            assertEquals(2, ffp.getRestHandlers(null, null, null, null, null, null, null).size());
-            assertEquals(2, ffp.getActions().size());
-=======
             assertEquals(3, ffp.createComponents(client, null, threadPool, null, null, null, null, null, null, null, null).size());
             assertEquals(3, ffp.getRestHandlers(null, null, null, null, null, null, null).size());
             assertEquals(3, ffp.getActions().size());
->>>>>>> 57869ff (Added javadoc)
             assertEquals(1, ffp.getExecutorBuilders(settings).size());
             assertEquals(1, ffp.getSettings().size());
         }

src/test/java/org/opensearch/flowframework/rest/RestSearchWorkflowActionTests.java

@@ -10,35 +10,42 @@
 
 import org.opensearch.client.node.NodeClient;
 import org.opensearch.core.common.bytes.BytesArray;
+import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.core.xcontent.XContentParseException;
+import org.opensearch.flowframework.common.FlowFrameworkFeatureEnabledSetting;
 import org.opensearch.rest.RestHandler.Route;
 import org.opensearch.rest.RestRequest;
 import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.test.rest.FakeRestChannel;
 import org.opensearch.test.rest.FakeRestRequest;
 
 import java.util.List;
 import java.util.Locale;
 
 import static org.opensearch.flowframework.common.CommonValue.WORKFLOW_URI;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 public class RestSearchWorkflowActionTests extends OpenSearchTestCase {
     private RestSearchWorkflowAction restSearchWorkflowAction;
     private String searchPath;
     private NodeClient nodeClient;
+    private FlowFrameworkFeatureEnabledSetting flowFrameworkFeatureEnabledSetting;
 
     @Override
     public void setUp() throws Exception {
         super.setUp();
 
         this.searchPath = String.format(Locale.ROOT, "%s/%s", WORKFLOW_URI, "_search");
-        this.restSearchWorkflowAction = new RestSearchWorkflowAction();
+        flowFrameworkFeatureEnabledSetting = mock(FlowFrameworkFeatureEnabledSetting.class);
+        when(flowFrameworkFeatureEnabledSetting.isFlowFrameworkEnabled()).thenReturn(true);
+        this.restSearchWorkflowAction = new RestSearchWorkflowAction(flowFrameworkFeatureEnabledSetting);
         this.nodeClient = mock(NodeClient.class);
     }
 
     public void testConstructor() {
-        RestSearchWorkflowAction searchWorkflowAction = new RestSearchWorkflowAction();
+        RestSearchWorkflowAction searchWorkflowAction = new RestSearchWorkflowAction(flowFrameworkFeatureEnabledSetting);
         assertNotNull(searchWorkflowAction);
     }
 
@@ -69,4 +76,15 @@ public void testInvalidSearchRequest() {
         });
         assertEquals("unknown named object category [org.opensearch.index.query.QueryBuilder]", ex.getMessage());
     }
+
+    public void testFeatureFlagNotEnabled() throws Exception {
+        when(flowFrameworkFeatureEnabledSetting.isFlowFrameworkEnabled()).thenReturn(false);
+        RestRequest request = new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.POST)
+                .withPath(this.searchPath)
+                .build();
+        FakeRestChannel channel = new FakeRestChannel(request, false, 1);
+        restSearchWorkflowAction.handleRequest(request, channel, nodeClient);
+        assertEquals(RestStatus.FORBIDDEN, channel.capturedResponse().status());
+        assertTrue(channel.capturedResponse().content().utf8ToString().contains("This API is disabled."));
+    }
 }