From 0c57adda04d2d0266d83a6370788a200c0069289 Mon Sep 17 00:00:00 2001
From: Joshua Palis <jpalis@amazon.com>
Date: Fri, 22 Mar 2024 19:14:11 +0000
Subject: [PATCH] Fixing default semantic use case integration test

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

diff --git a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
index 4acc80a0a..423d0ae8c 100644
--- a/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
+++ b/src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java
@@ -33,6 +33,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;
@@ -706,4 +708,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,
+            "",
+            ImmutableList.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 fb0ee879b..2d099a050 100644
--- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
+++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
@@ -438,7 +438,6 @@ public void testDefaultCohereUseCase() throws Exception {
     }
 
     public void testDefaultSemanticSearchUseCaseWithFailureExpected() throws Exception {
-
         // Hit Create Workflow API with original template
         Response response = createWorkflowWithUseCase(client(), "semantic_search");
         assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
@@ -455,17 +454,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"
+                )
+            );
+        }
     }
 
 }