From 14d370b167222503ffff030ac90da77ca09fb6e3 Mon Sep 17 00:00:00 2001
From: Joshua Palis <jpalis@amazon.com>
Date: Fri, 22 Mar 2024 19:55:27 +0000
Subject: [PATCH] backport

Signed-off-by: Joshua Palis <jpalis@amazon.com>
---
 .../FlowFrameworkRestTestCase.java            | 20 +++++++++++++++++
 .../rest/FlowFrameworkRestApiIT.java          | 22 ++++++++++++-------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
index c0a8975b8..4fe726ef8 100644
--- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
+++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
@@ -28,6 +28,8 @@
 import org.opensearch.common.settings.Settings;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.common.xcontent.LoggingDeprecationHandler;
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.commons.rest.SecureRestClientBuilder;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.DeprecationHandler;
@@ -681,4 +683,22 @@ protected GetPipelineResponse getPipelines() throws IOException {
             return GetPipelineResponse.fromXContent(parser);
         }
     }
+
+    @SuppressWarnings("unchecked")
+    protected List<String> catPlugins() throws IOException {
+        Response response = TestHelpers.makeRequest(
+            client(),
+            "GET",
+            "_cat/plugins?s=component&h=name,component,version,description&format=json",
+            null,
+            "",
+            List.of(new BasicHeader(HttpHeaders.USER_AGENT, ""))
+        );
+        List<Object> pluginsList = JsonXContent.jsonXContent.createParser(
+            NamedXContentRegistry.EMPTY,
+            LoggingDeprecationHandler.INSTANCE,
+            response.getEntity().getContent()
+        ).list();
+        return pluginsList.stream().map(o -> ((Map<String, Object>) o).get("component").toString()).collect(Collectors.toList());
+    }
 }
diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
index 65c3cd6d5..a8f786abc 100644
--- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
+++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
@@ -457,16 +457,22 @@ public void testDefaultSemanticSearchUseCaseWithFailureExpected() throws Excepti
             response = provisionWorkflow(client(), workflowId);
         }
 
-        // expecting a failure since there is no neural-search plugin in cluster to provide text-embedding processor
         assertEquals(RestStatus.OK, TestHelpers.restStatus(response));
-        getAndAssertWorkflowStatus(client(), workflowId, State.FAILED, ProvisioningProgress.FAILED);
 
-        String error = getAndWorkflowStatusError(client(), workflowId);
-        assertTrue(
-            error.contains(
-                "org.opensearch.flowframework.exception.WorkflowStepException during step create_ingest_pipeline, restStatus: BAD_REQUEST"
-            )
-        );
+        // Distribution build contains all plugins, checking if plugins are part of the integration test cluster
+        List<String> plugins = catPlugins();
+        if (plugins.contains("opensearch-knn") && plugins.contains("neural-search")) {
+            getAndAssertWorkflowStatus(client(), workflowId, State.PROVISIONING, ProvisioningProgress.IN_PROGRESS);
+        } else {
+            // expecting a failure since there is no neural-search plugin in cluster to provide text-embedding processor
+            getAndAssertWorkflowStatus(client(), workflowId, State.FAILED, ProvisioningProgress.FAILED);
+            String error = getAndWorkflowStatusError(client(), workflowId);
+            assertTrue(
+                error.contains(
+                    "org.opensearch.flowframework.exception.WorkflowStepException during step create_ingest_pipeline, restStatus: BAD_REQUEST"
+                )
+            );
+        }
     }
 
     public void testAllDefaultUseCasesCreation() throws Exception {