From 50b44eb8c7ed9c0735c0012ffd7791c0804cf2c5 Mon Sep 17 00:00:00 2001
From: Amit Galitzky <amgalitz@amazon.com>
Date: Fri, 7 Jun 2024 16:57:33 -0700
Subject: [PATCH] changing file structure for bwc

Signed-off-by: Amit Galitzky <amgalitz@amazon.com>
---
 build.gradle                                  | 24 ++++++++++++++++++-
 .../rest/FlowFrameworkRestApiIT.java          |  6 ++++-
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/build.gradle b/build.gradle
index 19287a4b4..698ae8a56 100644
--- a/build.gradle
+++ b/build.gradle
@@ -494,7 +494,29 @@ List<Provider<RegularFile>> plugins = [
                 return new RegularFile() {
                     @Override
                     File getAsFile() {
-                        return configurations.zipArchive.asFileTree.getFiles()
+                        return configurations.zipArchive.asFileTree.matching{include "**/opensearch-ml-plugin-${opensearch_build}.zip"}.getSingleFile()
+                    }
+                }
+            }
+        }),
+        provider(new Callable<RegularFile>(){
+            @Override
+            RegularFile call() throws Exception {
+                return new RegularFile() {
+                    @Override
+                    File getAsFile() {
+                        return configurations.zipArchive.asFileTree.matching{include "**/opensearch-knn-${opensearch_build}.zip"}.getSingleFile()
+                    }
+                }
+            }
+        }),
+        provider(new Callable<RegularFile>(){
+            @Override
+            RegularFile call() throws Exception {
+                return new RegularFile() {
+                    @Override
+                    File getAsFile() {
+                        return configurations.zipArchive.asFileTree.matching{include "**/neural-search-${opensearch_build}.zip"}.getSingleFile()
                     }
                 }
             }
diff --git a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
index 1ddf67c2a..6dfad72f1 100644
--- a/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
+++ b/src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java
@@ -525,7 +525,11 @@ public void testAllDefaultUseCasesCreation() throws Exception {
     }
 
     public void testSemanticSearchWithLocalModelEndToEnd() throws Exception {
-
+        // Checking if plugins are part of the integration test cluster so we can continue with this test
+        List<String> plugins = catPlugins();
+        if (!plugins.contains("opensearch-knn") && plugins.contains("opensearch-neural-search")) {
+            return;
+        }
         Map<String, Object> defaults = new HashMap<>();
         defaults.put("register_local_pretrained_model.name", "huggingface/sentence-transformers/all-MiniLM-L6-v2");
         defaults.put("register_local_pretrained_model.version", "1.0.1");