From c75fc5092bfd74087f0ee44a9fc15994d7db05fc Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Thu, 1 Jun 2023 15:03:54 -0700
Subject: [PATCH] Adding model group id to integ tests (#189) (#190)

* Fix Model group id in integ tests

Signed-off-by: Martin Gaievski <gaievski@amazon.com>

---------

Signed-off-by: Martin Gaievski <gaievski@amazon.com>
(cherry picked from commit 5228fc7a72b24543b6fd7c087cc062623ee00009)

Co-authored-by: Martin Gaievski <gaievski@amazon.com>
---
 .../common/BaseNeuralSearchIT.java            | 26 +++++++++++++++++++
 .../CreateModelGroupRequestBody.json          |  4 +++
 .../processor/UploadModelRequestBody.json     |  1 +
 3 files changed, 31 insertions(+)
 create mode 100644 src/test/resources/processor/CreateModelGroupRequestBody.json

diff --git a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
index 7f83e554b..0b5384b77 100644
--- a/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
+++ b/src/test/java/org/opensearch/neuralsearch/common/BaseNeuralSearchIT.java
@@ -9,6 +9,7 @@
 import static org.opensearch.neuralsearch.common.VectorUtil.vectorAsListToArray;
 
 import java.io.IOException;
+import java.net.URISyntaxException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.Collections;
@@ -88,6 +89,9 @@ protected void updateClusterSettings(String settingKey, Object value) {
     }
 
     protected String uploadModel(String requestBody) throws Exception {
+        String modelGroupId = registerModelGroup();
+        // model group id is dynamically generated, we need to update model update request body after group is registered
+        requestBody = requestBody.replace("<MODEL_GROUP_ID>", modelGroupId);
         Response uploadResponse = makeRequest(
             client(),
             "POST",
@@ -492,4 +496,26 @@ protected static class KNNFieldConfig {
         private final Integer dimension;
         private final SpaceType spaceType;
     }
+
+    private String registerModelGroup() throws IOException, URISyntaxException {
+        String modelGroupRegisterRequestBody = Files.readString(
+            Path.of(classLoader.getResource("processor/CreateModelGroupRequestBody.json").toURI())
+        );
+        Response modelGroupResponse = makeRequest(
+            client(),
+            "POST",
+            "/_plugins/_ml/model_groups/_register",
+            null,
+            toHttpEntity(modelGroupRegisterRequestBody),
+            ImmutableList.of(new BasicHeader(HttpHeaders.USER_AGENT, "Kibana"))
+        );
+        Map<String, Object> modelGroupResJson = XContentHelper.convertToMap(
+            XContentFactory.xContent(XContentType.JSON),
+            EntityUtils.toString(modelGroupResponse.getEntity()),
+            false
+        );
+        String modelGroupId = modelGroupResJson.get("model_group_id").toString();
+        assertNotNull(modelGroupId);
+        return modelGroupId;
+    }
 }
diff --git a/src/test/resources/processor/CreateModelGroupRequestBody.json b/src/test/resources/processor/CreateModelGroupRequestBody.json
new file mode 100644
index 000000000..2fddae02e
--- /dev/null
+++ b/src/test/resources/processor/CreateModelGroupRequestBody.json
@@ -0,0 +1,4 @@
+{
+  "name": "test_model_group_public",
+  "description": "This is a public model group"
+}
diff --git a/src/test/resources/processor/UploadModelRequestBody.json b/src/test/resources/processor/UploadModelRequestBody.json
index 9fc53f3b9..95f9c9cb5 100644
--- a/src/test/resources/processor/UploadModelRequestBody.json
+++ b/src/test/resources/processor/UploadModelRequestBody.json
@@ -4,6 +4,7 @@
   "model_format": "TORCH_SCRIPT",
   "model_task_type": "text_embedding",
   "model_content_hash_value": "e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021",
+  "model_group_id": "<MODEL_GROUP_ID>",
   "model_config": {
     "model_type": "bert",
     "embedding_dimension": 768,